<commit_before>use http::{is_valid_url, resolve_url, retrieve_asset};
use std::default::Default;
use std::io;
use utils::data_to_dataurl;

use html5ever::parse_document;
use html5ever::rcdom::{Handle, NodeData, RcDom};
use html5ever::serialize::{serialize, SerializeOpts};
use html5ever::tendril::TendrilSink;

enum NodeMatch {
    Icon,
    Image,
    Source,
    StyleSheet,
    Anchor,
    Script,
    Form,
    IFrame,
    Other,
}

const TRANSPARENT_PIXEL: &str = "data:image/png;base64,\
iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=";

const JS_DOM_EVENT_ATTRS: [&str; 21] = [
    // Input
    "onfocus",
    "onblur",
    "onselect",
    "onchange",
    "onsubmit",
    "onreset",
    "onkeydown",
    "onkeypress",
    "onkeyup",
    // Mouse
    "onmouseover",
    "onmouseout",
    "onmousedown",
    "onmouseup",
    "onmousemove",
    // Click
    "onclick",
    "ondblclick",
    // Load
    "onload",
    "onunload",
    "onabort",
    "onerror",
    "onresize",
];

fn get_parent_node_name(node: &Handle) -> String {
    let parent = node.parent.take().clone();
    let parent_node = parent.and_then(|node| node.upgrade()).unwrap();

    match &parent_node.data {
        NodeData::Document => {"".to_string()}
        NodeData::Doctype { .. } => {"".to_string()}
        NodeData::Text { .. } => {"".to_string()}
        NodeData::Comment { .. } => {"".to_string()}
        NodeData::Element {
            ref name,
            attrs: _,
            ..
        } => {
            name.local.as_ref().to_string()
        }
        NodeData::ProcessingInstruction { .. } => unreachable!()
    }
}

pub fn walk_and_embed_assets(
    url: &str,
    node: &Handle,
    opt_no_js: bool,
    opt_no_images: bool,
    opt_user_agent: &str,
) {
    match node.data {
        NodeData::Document => {
            // Dig deeper
            for child in node.children.borrow().iter() {
                walk_and_embed_assets(
                    &url, child,
                    opt_no_js,
                    opt_no_images,
                    opt_user_agent,
                );
            }
        }

        NodeData::Doctype { .. } => {}

        NodeData::Text { .. } => {}

        NodeData::Comment { .. } => {
            // Note: in case of opt_no_js being set to true, there's no need to worry about
            //       getting rid of comments that may contain scripts, e.g. <!--[if IE]><script>...
            //       since that's not part of W3C standard and therefore gets ignored
            //       by browsers other than IE [5, 9]
        }

        NodeData::Element {
            ref name,
            ref attrs,
            ..
        } => {
            let attrs_mut = &mut attrs.borrow_mut();
            let mut found = NodeMatch::Other;

            match name.local.as_ref() {
                "link" => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "rel" {
                            if is_icon(&attr.value.to_string()) {
                                found = NodeMatch::Icon;
                                break;
                            } else if attr.value.to_string() == "stylesheet" {
                                found = NodeMatch::StyleSheet;
                                break;
                            }
                        }
                    }
                }
                "img" => { found = NodeMatch::Image; }
                "source" => { found = NodeMatch::Source; }
                "a" => { found = NodeMatch::Anchor; }
                "script" => { found = NodeMatch::Script; }
                "form" => { found = NodeMatch::Form; }
                "iframe" => { found = NodeMatch::IFrame; }
                _ => {}
            }

            match found {
                NodeMatch::Icon => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "href" {
                            let href_full_url = resolve_url(&url, &attr.value.to_string());
                            let favicon_datauri = retrieve_asset(
                                &href_full_url.unwrap(),
                                true,
                                "",
                                opt_user_agent,
                            );
                            attr.value.clear();
                            attr.value.push_slice(favicon_datauri.unwrap().as_str());
                        }
                    }
                }
                NodeMatch::Image => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "src" {
                            if opt_no_images {
                                attr.value.clear();
                                attr.value.push_slice(TRANSPARENT_PIXEL);
                            } else {
                                let src_full_url = resolve_url(&url, &attr.value.to_string());
                                let img_datauri = retrieve_asset(
                                    &src_full_url.unwrap(),
                                    true,
                                    "",
                                    opt_user_agent,
                                );
                                attr.value.clear();
                                attr.value.push_slice(img_datauri.unwrap().as_str());
                            }
                        }
                    }
                }
                NodeMatch::Source => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "srcset" {
                            if get_parent_node_name(&node) == "picture" {
                                if opt_no_images {
                                    attr.value.clear();
                                    attr.value.push_slice(TRANSPARENT_PIXEL);
                                } else {
                                    let srcset_full_url = resolve_url(&url, &attr.value.to_string());
                                    let source_datauri = retrieve_asset(
                                        &srcset_full_url.unwrap(),
                                        true,
                                        "",
                                        opt_user_agent,
                                    );
                                    attr.value.clear();
                                    attr.value.push_slice(source_datauri.unwrap().as_str());
                                }
                            }
                        }
                    }
                }
                NodeMatch::Anchor => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "href" {
                            // Don't touch email links or hrefs which begin with a hash sign
                            if attr.value.starts_with('#') || attr.value.starts_with("mailto:") {
                                continue;
                            }

                            let href_full_url = resolve_url(&url, &attr.value.to_string());
                            attr.value.clear();
                            attr.value.push_slice(href_full_url.unwrap().as_str());
                        }
                    }
                }
                NodeMatch::StyleSheet => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "href" {
                            let href_full_url = resolve_url(&url, &attr.value.to_string());
                            let css_datauri = retrieve_asset(
                                &href_full_url.unwrap(),
                                true,
                                "text/css",
                                opt_user_agent,
                            );
                            attr.value.clear();
                            attr.value.push_slice(css_datauri.unwrap().as_str());
                        }
                    }
                }
                NodeMatch::Script => {
                    if opt_no_js {
                        // Get rid of src and inner content of SCRIPT tags
                        for attr in attrs_mut.iter_mut() {
                            if &attr.name.local == "src" {
                                attr.value.clear();
                            }
                        }
                        node.children.borrow_mut().clear();
                    } else {
                        for attr in attrs_mut.iter_mut() {
                            if &attr.name.local == "src" {
                                let src_full_url = resolve_url(&url, &attr.value.to_string());
                                let js_datauri = retrieve_asset(
                                    &src_full_url.unwrap(),
                                    true,
                                    "application/javascript",
                                    opt_user_agent,
                                );
                                attr.value.clear();
                                attr.value.push_slice(js_datauri.unwrap().as_str());
                            }
                        }
                    }
                }
                NodeMatch::Form => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "action" {
                            // Do not touch action props which are set to a URL
                            if is_valid_url(&attr.value) {
                                continue;
                            }

                            let href_full_url = resolve_url(&url, &attr.value.to_string());
                            attr.value.clear();
                            attr.value.push_slice(href_full_url.unwrap().as_str());
                        }
                    }
                }
                NodeMatch::IFrame => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "src" {
                            let src_full_url = resolve_url(&url, &attr.value.to_string()).unwrap();
                            let iframe_data = retrieve_asset(
                                &src_full_url,
                                false,
                                "text/html",
                                opt_user_agent,
                            );
                            let dom = html_to_dom(&iframe_data.unwrap());
                            walk_and_embed_assets(
                                &src_full_url,
                                &dom.document,
                                opt_no_js,
                                opt_no_images,
                                opt_user_agent,
                            );
                            let mut buf: Vec<u8> = Vec::new();
                            serialize(&mut buf, &dom.document, SerializeOpts::default()).unwrap();
                            let iframe_datauri = data_to_dataurl("text/html", &buf);
                            attr.value.clear();
                            attr.value.push_slice(iframe_datauri.as_str());
                        }
                    }
                }
                NodeMatch::Other => {}
            }

            if opt_no_js {
                // Get rid of JS event attributes
                for attr in attrs_mut.iter_mut() {
                    if JS_DOM_EVENT_ATTRS.contains(&attr.name.local.to_lowercase().as_str()) {
                        attr.value.clear();
                    }
                }
            }

            // Dig deeper
            for child in node.children.borrow().iter() {
                walk_and_embed_assets(
                    &url,
                    child,
                    opt_no_js,
                    opt_no_images,
                    opt_user_agent,
                );
            }
        }

        NodeData::ProcessingInstruction { .. } => unreachable!()
    }
}

pub fn html_to_dom(data: &str) -> html5ever::rcdom::RcDom {
    parse_document(RcDom::default(), Default::default())
        .from_utf8()
        .read_from(&mut data.as_bytes())
        .unwrap()
}

pub fn print_dom(handle: &Handle) {
    serialize(&mut io::stdout(), handle, SerializeOpts::default()).unwrap();
}

fn is_icon(attr_value: &str) -> bool {
    attr_value == "icon"
        || attr_value == "shortcut icon"
        || attr_value == "mask-icon"
        || attr_value == "apple-touch-icon"
}

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

    #[test]
    fn test_is_icon() {
        assert_eq!(is_icon("icon"), true);
        assert_eq!(is_icon("stylesheet"), false);
    }
}
<commit_msg>Ignore icons if told to avoid images<commit_after>use http::{is_valid_url, resolve_url, retrieve_asset};
use std::default::Default;
use std::io;
use utils::data_to_dataurl;

use html5ever::parse_document;
use html5ever::rcdom::{Handle, NodeData, RcDom};
use html5ever::serialize::{serialize, SerializeOpts};
use html5ever::tendril::TendrilSink;

enum NodeMatch {
    Icon,
    Image,
    Source,
    StyleSheet,
    Anchor,
    Script,
    Form,
    IFrame,
    Other,
}

const TRANSPARENT_PIXEL: &str = "data:image/png;base64,\
iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=";

const JS_DOM_EVENT_ATTRS: [&str; 21] = [
    // Input
    "onfocus",
    "onblur",
    "onselect",
    "onchange",
    "onsubmit",
    "onreset",
    "onkeydown",
    "onkeypress",
    "onkeyup",
    // Mouse
    "onmouseover",
    "onmouseout",
    "onmousedown",
    "onmouseup",
    "onmousemove",
    // Click
    "onclick",
    "ondblclick",
    // Load
    "onload",
    "onunload",
    "onabort",
    "onerror",
    "onresize",
];

fn get_parent_node_name(node: &Handle) -> String {
    let parent = node.parent.take().clone();
    let parent_node = parent.and_then(|node| node.upgrade()).unwrap();

    match &parent_node.data {
        NodeData::Document => {"".to_string()}
        NodeData::Doctype { .. } => {"".to_string()}
        NodeData::Text { .. } => {"".to_string()}
        NodeData::Comment { .. } => {"".to_string()}
        NodeData::Element {
            ref name,
            attrs: _,
            ..
        } => {
            name.local.as_ref().to_string()
        }
        NodeData::ProcessingInstruction { .. } => unreachable!()
    }
}

pub fn walk_and_embed_assets(
    url: &str,
    node: &Handle,
    opt_no_js: bool,
    opt_no_images: bool,
    opt_user_agent: &str,
) {
    match node.data {
        NodeData::Document => {
            // Dig deeper
            for child in node.children.borrow().iter() {
                walk_and_embed_assets(
                    &url, child,
                    opt_no_js,
                    opt_no_images,
                    opt_user_agent,
                );
            }
        }

        NodeData::Doctype { .. } => {}

        NodeData::Text { .. } => {}

        NodeData::Comment { .. } => {
            // Note: in case of opt_no_js being set to true, there's no need to worry about
            //       getting rid of comments that may contain scripts, e.g. <!--[if IE]><script>...
            //       since that's not part of W3C standard and therefore gets ignored
            //       by browsers other than IE [5, 9]
        }

        NodeData::Element {
            ref name,
            ref attrs,
            ..
        } => {
            let attrs_mut = &mut attrs.borrow_mut();
            let mut found = NodeMatch::Other;

            match name.local.as_ref() {
                "link" => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "rel" {
                            if is_icon(&attr.value.to_string()) {
                                found = NodeMatch::Icon;
                                break;
                            } else if attr.value.to_string() == "stylesheet" {
                                found = NodeMatch::StyleSheet;
                                break;
                            }
                        }
                    }
                }
                "img" => { found = NodeMatch::Image; }
                "source" => { found = NodeMatch::Source; }
                "a" => { found = NodeMatch::Anchor; }
                "script" => { found = NodeMatch::Script; }
                "form" => { found = NodeMatch::Form; }
                "iframe" => { found = NodeMatch::IFrame; }
                _ => {}
            }

            match found {
                NodeMatch::Icon => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "href" {
                            if opt_no_images {
                                attr.value.clear();
                                attr.value.push_slice(TRANSPARENT_PIXEL);
                            } else {
                                let href_full_url = resolve_url(&url, &attr.value.to_string());
                                let favicon_datauri = retrieve_asset(
                                    &href_full_url.unwrap(),
                                    true,
                                    "",
                                    opt_user_agent,
                                );
                                attr.value.clear();
                                attr.value.push_slice(favicon_datauri.unwrap().as_str());
                            }
                        }
                    }
                }
                NodeMatch::Image => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "src" {
                            if opt_no_images {
                                attr.value.clear();
                                attr.value.push_slice(TRANSPARENT_PIXEL);
                            } else {
                                let src_full_url = resolve_url(&url, &attr.value.to_string());
                                let img_datauri = retrieve_asset(
                                    &src_full_url.unwrap(),
                                    true,
                                    "",
                                    opt_user_agent,
                                );
                                attr.value.clear();
                                attr.value.push_slice(img_datauri.unwrap().as_str());
                            }
                        }
                    }
                }
                NodeMatch::Source => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "srcset" {
                            if get_parent_node_name(&node) == "picture" {
                                if opt_no_images {
                                    attr.value.clear();
                                    attr.value.push_slice(TRANSPARENT_PIXEL);
                                } else {
                                    let srcset_full_url = resolve_url(&url, &attr.value.to_string());
                                    let source_datauri = retrieve_asset(
                                        &srcset_full_url.unwrap(),
                                        true,
                                        "",
                                        opt_user_agent,
                                    );
                                    attr.value.clear();
                                    attr.value.push_slice(source_datauri.unwrap().as_str());
                                }
                            }
                        }
                    }
                }
                NodeMatch::Anchor => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "href" {
                            // Don't touch email links or hrefs which begin with a hash sign
                            if attr.value.starts_with('#') || attr.value.starts_with("mailto:") {
                                continue;
                            }

                            let href_full_url = resolve_url(&url, &attr.value.to_string());
                            attr.value.clear();
                            attr.value.push_slice(href_full_url.unwrap().as_str());
                        }
                    }
                }
                NodeMatch::StyleSheet => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "href" {
                            let href_full_url = resolve_url(&url, &attr.value.to_string());
                            let css_datauri = retrieve_asset(
                                &href_full_url.unwrap(),
                                true,
                                "text/css",
                                opt_user_agent,
                            );
                            attr.value.clear();
                            attr.value.push_slice(css_datauri.unwrap().as_str());
                        }
                    }
                }
                NodeMatch::Script => {
                    if opt_no_js {
                        // Get rid of src and inner content of SCRIPT tags
                        for attr in attrs_mut.iter_mut() {
                            if &attr.name.local == "src" {
                                attr.value.clear();
                            }
                        }
                        node.children.borrow_mut().clear();
                    } else {
                        for attr in attrs_mut.iter_mut() {
                            if &attr.name.local == "src" {
                                let src_full_url = resolve_url(&url, &attr.value.to_string());
                                let js_datauri = retrieve_asset(
                                    &src_full_url.unwrap(),
                                    true,
                                    "application/javascript",
                                    opt_user_agent,
                                );
                                attr.value.clear();
                                attr.value.push_slice(js_datauri.unwrap().as_str());
                            }
                        }
                    }
                }
                NodeMatch::Form => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "action" {
                            // Do not touch action props which are set to a URL
                            if is_valid_url(&attr.value) {
                                continue;
                            }

                            let href_full_url = resolve_url(&url, &attr.value.to_string());
                            attr.value.clear();
                            attr.value.push_slice(href_full_url.unwrap().as_str());
                        }
                    }
                }
                NodeMatch::IFrame => {
                    for attr in attrs_mut.iter_mut() {
                        if &attr.name.local == "src" {
                            let src_full_url = resolve_url(&url, &attr.value.to_string()).unwrap();
                            let iframe_data = retrieve_asset(
                                &src_full_url,
                                false,
                                "text/html",
                                opt_user_agent,
                            );
                            let dom = html_to_dom(&iframe_data.unwrap());
                            walk_and_embed_assets(
                                &src_full_url,
                                &dom.document,
                                opt_no_js,
                                opt_no_images,
                                opt_user_agent,
                            );
                            let mut buf: Vec<u8> = Vec::new();
                            serialize(&mut buf, &dom.document, SerializeOpts::default()).unwrap();
                            let iframe_datauri = data_to_dataurl("text/html", &buf);
                            attr.value.clear();
                            attr.value.push_slice(iframe_datauri.as_str());
                        }
                    }
                }
                NodeMatch::Other => {}
            }

            if opt_no_js {
                // Get rid of JS event attributes
                for attr in attrs_mut.iter_mut() {
                    if JS_DOM_EVENT_ATTRS.contains(&attr.name.local.to_lowercase().as_str()) {
                        attr.value.clear();
                    }
                }
            }

            // Dig deeper
            for child in node.children.borrow().iter() {
                walk_and_embed_assets(
                    &url,
                    child,
                    opt_no_js,
                    opt_no_images,
                    opt_user_agent,
                );
            }
        }

        NodeData::ProcessingInstruction { .. } => unreachable!()
    }
}

pub fn html_to_dom(data: &str) -> html5ever::rcdom::RcDom {
    parse_document(RcDom::default(), Default::default())
        .from_utf8()
        .read_from(&mut data.as_bytes())
        .unwrap()
}

pub fn print_dom(handle: &Handle) {
    serialize(&mut io::stdout(), handle, SerializeOpts::default()).unwrap();
}

fn is_icon(attr_value: &str) -> bool {
    attr_value == "icon"
        || attr_value == "shortcut icon"
        || attr_value == "mask-icon"
        || attr_value == "apple-touch-icon"
}

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

    #[test]
    fn test_is_icon() {
        assert_eq!(is_icon("icon"), true);
        assert_eq!(is_icon("stylesheet"), false);
    }
}
<|endoftext|>


<commit_before>#![allow(unsafe_code)]

use std::{
    alloc::{alloc, dealloc, Layout},
    convert::TryFrom,
    fmt,
    hash::{Hash, Hasher},
    iter::FromIterator,
    mem::size_of,
    ops::{Deref, DerefMut},
    sync::atomic::{AtomicUsize, Ordering},
};

const SZ: usize = size_of::<usize>();
const CUTOFF: usize = SZ - 1;

/// A buffer that may either be inline or remote and protected
/// by an Arc. The inner buffer is guaranteed to be aligned to
/// 8 byte boundaries.
#[repr(align(8))]
pub struct IVec([u8; SZ]);

impl Clone for IVec {
    fn clone(&self) -> IVec {
        if !self.is_inline() {
            self.deref_header().rc.fetch_add(1, Ordering::Relaxed);
        }
        IVec(self.0)
    }
}

impl Drop for IVec {
    fn drop(&mut self) {
        if !self.is_inline() {
            let rc = self.deref_header().rc.fetch_sub(1, Ordering::Release) - 1;

            if rc == 0 {
                let layout = Layout::from_size_align(
                    self.deref_header().len + size_of::<RemoteHeader>(),
                    8,
                )
                .unwrap();

                std::sync::atomic::fence(Ordering::Acquire);

                unsafe {
                    dealloc(self.remote_ptr() as *mut u8, layout);
                }
            }
        }
    }
}

struct RemoteHeader {
    rc: AtomicUsize,
    len: usize,
}

impl Deref for IVec {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        if self.is_inline() {
            &self.0[..self.inline_len()]
        } else {
            unsafe {
                let data_ptr = self.remote_ptr().add(size_of::<RemoteHeader>());
                let len = self.deref_header().len;
                std::slice::from_raw_parts(data_ptr, len)
            }
        }
    }
}

impl AsRef<[u8]> for IVec {
    #[inline]
    fn as_ref(&self) -> &[u8] {
        self
    }
}

impl DerefMut for IVec {
    #[inline]
    fn deref_mut(&mut self) -> &mut [u8] {
        let inline_len = self.inline_len();
        if self.is_inline() {
            &mut self.0[..inline_len]
        } else {
            self.make_mut();
            unsafe {
                let data_ptr = self.remote_ptr().add(size_of::<RemoteHeader>());
                let len = self.deref_header().len;
                std::slice::from_raw_parts_mut(data_ptr as *mut u8, len)
            }
        }
    }
}

impl AsMut<[u8]> for IVec {
    #[inline]
    fn as_mut(&mut self) -> &mut [u8] {
        self.deref_mut()
    }
}

impl Default for IVec {
    fn default() -> Self {
        Self::from(&[])
    }
}

impl Hash for IVec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.deref().hash(state);
    }
}

impl IVec {
    fn new(slice: &[u8]) -> Self {
        let mut data = [0_u8; SZ];
        if slice.len() <= CUTOFF {
            data[SZ - 1] = (u8::try_from(slice.len()).unwrap() << 1) | 1;
            data[..slice.len()].copy_from_slice(slice);
        } else {
            let layout = Layout::from_size_align(
                slice.len() + size_of::<RemoteHeader>(),
                8,
            )
            .unwrap();

            let header = RemoteHeader { rc: 1.into(), len: slice.len() };

            unsafe {
                let ptr = alloc(layout);

                std::ptr::write(ptr as *mut RemoteHeader, header);
                std::ptr::copy_nonoverlapping(
                    slice.as_ptr(),
                    ptr.add(size_of::<RemoteHeader>()),
                    slice.len(),
                );
                std::ptr::write_unaligned(data.as_mut_ptr() as _, ptr);
            }

            // assert that the bottom 3 bits are empty, as we expect
            // the buffer to always have an alignment of 8 (2 ^ 3).
            #[cfg(not(miri))]
            assert_eq!(data[SZ - 1] & 0b111, 0);
        }
        Self(data)
    }

    fn remote_ptr(&self) -> *const u8 {
        assert!(!self.is_inline());
        unsafe { std::ptr::read(self.0.as_ptr() as *const *const u8) }
    }

    fn deref_header(&self) -> &RemoteHeader {
        assert!(!self.is_inline());
        unsafe { &*(self.remote_ptr() as *mut RemoteHeader) }
    }

    #[cfg(miri)]
    fn inline_len(&self) -> usize {
        (self.trailer() >> 1) as usize
    }

    #[cfg(miri)]
    fn is_inline(&self) -> bool {
        self.trailer() & 1 == 1
    }

    #[cfg(miri)]
    fn trailer(&self) -> u8 {
        self.deref()[SZ - 1]
    }

    #[cfg(not(miri))]
    const fn inline_len(&self) -> usize {
        (self.trailer() >> 1) as usize
    }

    #[cfg(not(miri))]
    const fn is_inline(&self) -> bool {
        self.trailer() & 1 == 1
    }

    #[cfg(not(miri))]
    const fn trailer(&self) -> u8 {
        self.0[SZ - 1]
    }

    fn make_mut(&mut self) {
        assert!(!self.is_inline());
        if self.deref_header().rc.load(Ordering::Acquire) != 1 {
            *self = IVec::from(self.deref())
        }
    }
}

impl FromIterator<u8> for IVec {
    fn from_iter<T>(iter: T) -> Self
    where
        T: IntoIterator<Item = u8>,
    {
        let bs: Vec<u8> = iter.into_iter().collect();
        bs.into()
    }
}

impl From<&[u8]> for IVec {
    fn from(slice: &[u8]) -> Self {
        IVec::new(slice)
    }
}

impl From<&str> for IVec {
    fn from(s: &str) -> Self {
        Self::from(s.as_bytes())
    }
}

impl From<String> for IVec {
    fn from(s: String) -> Self {
        Self::from(s.as_bytes())
    }
}

impl From<&String> for IVec {
    fn from(s: &String) -> Self {
        Self::from(s.as_bytes())
    }
}

impl From<&IVec> for IVec {
    fn from(v: &Self) -> Self {
        v.clone()
    }
}

impl From<Vec<u8>> for IVec {
    fn from(v: Vec<u8>) -> Self {
        IVec::new(&v)
    }
}

impl From<Box<[u8]>> for IVec {
    fn from(v: Box<[u8]>) -> Self {
        IVec::new(&v)
    }
}

impl std::borrow::Borrow<[u8]> for IVec {
    fn borrow(&self) -> &[u8] {
        self.as_ref()
    }
}

impl std::borrow::Borrow<[u8]> for &IVec {
    fn borrow(&self) -> &[u8] {
        self.as_ref()
    }
}

macro_rules! from_array {
    ($($s:expr),*) => {
        $(
            impl From<&[u8; $s]> for IVec {
                fn from(v: &[u8; $s]) -> Self {
                    Self::from(&v[..])
                }
            }
        )*
    }
}

from_array!(
    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
    21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32
);

impl Ord for IVec {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.as_ref().cmp(other.as_ref())
    }
}

impl PartialOrd for IVec {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<T: AsRef<[u8]>> PartialEq<T> for IVec {
    fn eq(&self, other: &T) -> bool {
        self.as_ref() == other.as_ref()
    }
}

impl PartialEq<[u8]> for IVec {
    fn eq(&self, other: &[u8]) -> bool {
        self.as_ref() == other
    }
}

impl Eq for IVec {}

impl fmt::Debug for IVec {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_ref().fmt(f)
    }
}

#[cfg(test)]
mod qc {
    use super::IVec;

    #[test]
    fn ivec_usage() {
        let iv1 = IVec::from(vec![1, 2, 3]);
        assert_eq!(iv1, vec![1, 2, 3]);
        let iv2 = IVec::from(&[4; 128][..]);
        assert_eq!(iv2, vec![4; 128]);
    }

    #[test]
    fn boxed_slice_conversion() {
        let boite1: Box<[u8]> = Box::new([1, 2, 3]);
        let iv1: IVec = boite1.into();
        assert_eq!(iv1, vec![1, 2, 3]);
        let boite2: Box<[u8]> = Box::new([4; 128]);
        let iv2: IVec = boite2.into();
        assert_eq!(iv2, vec![4; 128]);
    }

    #[test]
    fn ivec_as_mut_identity() {
        let initial = &[1];
        let mut iv = IVec::from(initial);
        assert_eq!(&*initial, &*iv);
        assert_eq!(&*initial, &mut *iv);
        assert_eq!(&*initial, iv.as_mut());
    }

    fn prop_identity(ivec: &IVec) -> bool {
        let mut iv2 = ivec.clone();

        if iv2 != ivec {
            println!("expected clone to equal original");
            return false;
        }

        if *ivec != *iv2 {
            println!("expected AsMut to equal original");
            return false;
        }

        if *ivec != iv2.as_mut() {
            println!("expected AsMut to equal original");
            return false;
        }

        true
    }

    quickcheck::quickcheck! {
        #[cfg_attr(miri, ignore)]
        fn ivec(item: IVec) -> bool {
            prop_identity(&item)
        }
    }

    #[test]
    fn ivec_bug_00() {
        assert!(prop_identity(&IVec::new(&[
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        ])));
    }
}
<commit_msg>Use const generics for IVec::from array<commit_after>#![allow(unsafe_code)]

use std::{
    alloc::{alloc, dealloc, Layout},
    convert::TryFrom,
    fmt,
    hash::{Hash, Hasher},
    iter::FromIterator,
    mem::size_of,
    ops::{Deref, DerefMut},
    sync::atomic::{AtomicUsize, Ordering},
};

const SZ: usize = size_of::<usize>();
const CUTOFF: usize = SZ - 1;

/// A buffer that may either be inline or remote and protected
/// by an Arc. The inner buffer is guaranteed to be aligned to
/// 8 byte boundaries.
#[repr(align(8))]
pub struct IVec([u8; SZ]);

impl Clone for IVec {
    fn clone(&self) -> IVec {
        if !self.is_inline() {
            self.deref_header().rc.fetch_add(1, Ordering::Relaxed);
        }
        IVec(self.0)
    }
}

impl Drop for IVec {
    fn drop(&mut self) {
        if !self.is_inline() {
            let rc = self.deref_header().rc.fetch_sub(1, Ordering::Release) - 1;

            if rc == 0 {
                let layout = Layout::from_size_align(
                    self.deref_header().len + size_of::<RemoteHeader>(),
                    8,
                )
                .unwrap();

                std::sync::atomic::fence(Ordering::Acquire);

                unsafe {
                    dealloc(self.remote_ptr() as *mut u8, layout);
                }
            }
        }
    }
}

struct RemoteHeader {
    rc: AtomicUsize,
    len: usize,
}

impl Deref for IVec {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        if self.is_inline() {
            &self.0[..self.inline_len()]
        } else {
            unsafe {
                let data_ptr = self.remote_ptr().add(size_of::<RemoteHeader>());
                let len = self.deref_header().len;
                std::slice::from_raw_parts(data_ptr, len)
            }
        }
    }
}

impl AsRef<[u8]> for IVec {
    #[inline]
    fn as_ref(&self) -> &[u8] {
        self
    }
}

impl DerefMut for IVec {
    #[inline]
    fn deref_mut(&mut self) -> &mut [u8] {
        let inline_len = self.inline_len();
        if self.is_inline() {
            &mut self.0[..inline_len]
        } else {
            self.make_mut();
            unsafe {
                let data_ptr = self.remote_ptr().add(size_of::<RemoteHeader>());
                let len = self.deref_header().len;
                std::slice::from_raw_parts_mut(data_ptr as *mut u8, len)
            }
        }
    }
}

impl AsMut<[u8]> for IVec {
    #[inline]
    fn as_mut(&mut self) -> &mut [u8] {
        self.deref_mut()
    }
}

impl Default for IVec {
    fn default() -> Self {
        Self::from(&[])
    }
}

impl Hash for IVec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.deref().hash(state);
    }
}

impl IVec {
    fn new(slice: &[u8]) -> Self {
        let mut data = [0_u8; SZ];
        if slice.len() <= CUTOFF {
            data[SZ - 1] = (u8::try_from(slice.len()).unwrap() << 1) | 1;
            data[..slice.len()].copy_from_slice(slice);
        } else {
            let layout = Layout::from_size_align(
                slice.len() + size_of::<RemoteHeader>(),
                8,
            )
            .unwrap();

            let header = RemoteHeader { rc: 1.into(), len: slice.len() };

            unsafe {
                let ptr = alloc(layout);

                std::ptr::write(ptr as *mut RemoteHeader, header);
                std::ptr::copy_nonoverlapping(
                    slice.as_ptr(),
                    ptr.add(size_of::<RemoteHeader>()),
                    slice.len(),
                );
                std::ptr::write_unaligned(data.as_mut_ptr() as _, ptr);
            }

            // assert that the bottom 3 bits are empty, as we expect
            // the buffer to always have an alignment of 8 (2 ^ 3).
            #[cfg(not(miri))]
            assert_eq!(data[SZ - 1] & 0b111, 0);
        }
        Self(data)
    }

    fn remote_ptr(&self) -> *const u8 {
        assert!(!self.is_inline());
        unsafe { std::ptr::read(self.0.as_ptr() as *const *const u8) }
    }

    fn deref_header(&self) -> &RemoteHeader {
        assert!(!self.is_inline());
        unsafe { &*(self.remote_ptr() as *mut RemoteHeader) }
    }

    #[cfg(miri)]
    fn inline_len(&self) -> usize {
        (self.trailer() >> 1) as usize
    }

    #[cfg(miri)]
    fn is_inline(&self) -> bool {
        self.trailer() & 1 == 1
    }

    #[cfg(miri)]
    fn trailer(&self) -> u8 {
        self.deref()[SZ - 1]
    }

    #[cfg(not(miri))]
    const fn inline_len(&self) -> usize {
        (self.trailer() >> 1) as usize
    }

    #[cfg(not(miri))]
    const fn is_inline(&self) -> bool {
        self.trailer() & 1 == 1
    }

    #[cfg(not(miri))]
    const fn trailer(&self) -> u8 {
        self.0[SZ - 1]
    }

    fn make_mut(&mut self) {
        assert!(!self.is_inline());
        if self.deref_header().rc.load(Ordering::Acquire) != 1 {
            *self = IVec::from(self.deref())
        }
    }
}

impl FromIterator<u8> for IVec {
    fn from_iter<T>(iter: T) -> Self
    where
        T: IntoIterator<Item = u8>,
    {
        let bs: Vec<u8> = iter.into_iter().collect();
        bs.into()
    }
}

impl From<&[u8]> for IVec {
    fn from(slice: &[u8]) -> Self {
        IVec::new(slice)
    }
}

impl From<&str> for IVec {
    fn from(s: &str) -> Self {
        Self::from(s.as_bytes())
    }
}

impl From<String> for IVec {
    fn from(s: String) -> Self {
        Self::from(s.as_bytes())
    }
}

impl From<&String> for IVec {
    fn from(s: &String) -> Self {
        Self::from(s.as_bytes())
    }
}

impl From<&IVec> for IVec {
    fn from(v: &Self) -> Self {
        v.clone()
    }
}

impl From<Vec<u8>> for IVec {
    fn from(v: Vec<u8>) -> Self {
        IVec::new(&v)
    }
}

impl From<Box<[u8]>> for IVec {
    fn from(v: Box<[u8]>) -> Self {
        IVec::new(&v)
    }
}

impl std::borrow::Borrow<[u8]> for IVec {
    fn borrow(&self) -> &[u8] {
        self.as_ref()
    }
}

impl std::borrow::Borrow<[u8]> for &IVec {
    fn borrow(&self) -> &[u8] {
        self.as_ref()
    }
}

impl<const N: usize> From<&[u8; N]> for IVec {
    fn from(v: &[u8; N]) -> Self {
        Self::from(&v[..])
    }
}

impl Ord for IVec {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.as_ref().cmp(other.as_ref())
    }
}

impl PartialOrd for IVec {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<T: AsRef<[u8]>> PartialEq<T> for IVec {
    fn eq(&self, other: &T) -> bool {
        self.as_ref() == other.as_ref()
    }
}

impl PartialEq<[u8]> for IVec {
    fn eq(&self, other: &[u8]) -> bool {
        self.as_ref() == other
    }
}

impl Eq for IVec {}

impl fmt::Debug for IVec {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_ref().fmt(f)
    }
}

#[cfg(test)]
mod qc {
    use super::IVec;

    #[test]
    fn ivec_usage() {
        let iv1 = IVec::from(vec![1, 2, 3]);
        assert_eq!(iv1, vec![1, 2, 3]);
        let iv2 = IVec::from(&[4; 128][..]);
        assert_eq!(iv2, vec![4; 128]);
    }

    #[test]
    fn boxed_slice_conversion() {
        let boite1: Box<[u8]> = Box::new([1, 2, 3]);
        let iv1: IVec = boite1.into();
        assert_eq!(iv1, vec![1, 2, 3]);
        let boite2: Box<[u8]> = Box::new([4; 128]);
        let iv2: IVec = boite2.into();
        assert_eq!(iv2, vec![4; 128]);
    }

    #[test]
    fn ivec_as_mut_identity() {
        let initial = &[1];
        let mut iv = IVec::from(initial);
        assert_eq!(&*initial, &*iv);
        assert_eq!(&*initial, &mut *iv);
        assert_eq!(&*initial, iv.as_mut());
    }

    fn prop_identity(ivec: &IVec) -> bool {
        let mut iv2 = ivec.clone();

        if iv2 != ivec {
            println!("expected clone to equal original");
            return false;
        }

        if *ivec != *iv2 {
            println!("expected AsMut to equal original");
            return false;
        }

        if *ivec != iv2.as_mut() {
            println!("expected AsMut to equal original");
            return false;
        }

        true
    }

    quickcheck::quickcheck! {
        #[cfg_attr(miri, ignore)]
        fn ivec(item: IVec) -> bool {
            prop_identity(&item)
        }
    }

    #[test]
    fn ivec_bug_00() {
        assert!(prop_identity(&IVec::new(&[
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        ])));
    }
}
<|endoftext|>


<commit_before>export prime, init, grow, get_at, iterable_prime, loopable_prime;

type prime = {
    mutable vec: [u64]
};

impl iterable_prime of iter::iterable<u64> for prime {
    fn iter(blk: fn(&&u64)) {
        let i = 0u;
        while true {
            blk(get_at_vec(self.vec, i));
            i += 1u;
        }
    }
}

impl loopable_prime for prime {
    fn loop(blk: fn(&&u64) -> bool) {
        let i = 0u;
        while blk(get_at_vec(self.vec, i)) {
            i += 1u;
        }
    }
}

fn init() -> prime {
    { mutable vec: [] }
}

fn grow_vec(&v: [u64], n: uint) {
    if n == 0u {
        ret;
    }
    let num = alt vec::last(v) {
      none       { v = [2u64];  grow_vec(v, n - 1u); ret }
      some(2u64) { v += [3u64]; grow_vec(v, n - 1u); ret }
      some(x)    { x + 2u64 }
    };

    let i = 0u;
    while i < n {
        for p in v {
            if p * p > num {
                v += [num];
                i += 1u;
                break;
            }
            if num % p == 0u64 {
                break;
            }
        }
        num += 2u;
    }
}

fn get_at_vec(&v: [u64], n: u64) -> u64 {
    if vec::len(v) <= n {
        grow_vec(v, n - vec::len(v) + 1u);
    }
    ret v[n];
}

fn grow(ps: prime, n: uint) {
    grow_vec(ps.vec, n);
}

fn get_at(&ps: prime, n: u64) -> u64 {
    if vec::len(ps.vec) <= n {
        grow(ps, n - vec::len(ps.vec));
    }
    ret ps.vec[n];
}
<commit_msg>Modify prime module's bug<commit_after>export prime, init, grow, get_at, iterable_prime, loopable_prime;

type prime = {
    mutable vec: [u64]
};

impl iterable_prime of iter::iterable<u64> for prime {
    fn iter(blk: fn(&&u64)) {
        let i = 0u;
        while true {
            blk(get_at_vec(self.vec, i));
            i += 1u;
        }
    }
}

impl loopable_prime for prime {
    fn loop(blk: fn(&&u64) -> bool) {
        let i = 0u;
        while blk(get_at_vec(self.vec, i)) {
            i += 1u;
        }
    }
}

fn init() -> prime {
    { mutable vec: [] }
}

fn grow_vec(&v: [u64], n: uint) {
    if n == 0u {
        ret;
    }
    let num = alt vec::last(v) {
      none       { v = [2u64];  grow_vec(v, n - 1u); ret }
      some(2u64) { v += [3u64]; grow_vec(v, n - 1u); ret }
      some(x)    { x + 2u64 }
    };

    let i = 0u;
    while i < n {
        for p in v {
            if p * p > num {
                v += [num];
                i += 1u;
                break;
            }
            if num % p == 0u64 {
                break;
            }
        }
        num += 2u;
    }
}

fn get_at_vec(&v: [u64], n: u64) -> u64 {
    if vec::len(v) <= n {
        grow_vec(v, n - vec::len(v) + 1u);
    }
    ret v[n];
}

fn grow(ps: prime, n: uint) {
    grow_vec(ps.vec, n);
}

fn get_at(&ps: prime, n: u64) -> u64 {
    ret get_at_vec(ps.vec, n);
}
<|endoftext|>


<commit_before>use std::{
    marker::PhantomData,
    hash::{Hash, Hasher},
};

use ra_db::{LocationIntener, FileId};
use ra_syntax::{TreeArc, SyntaxNode, SourceFile, AstNode, ast};
use ra_arena::{Arena, RawId, ArenaId, impl_arena_id};

use crate::{
    HirDatabase,
    Module,
};

#[derive(Debug, Default)]
pub struct HirInterner {
    macros: LocationIntener<MacroCallLoc, MacroCallId>,
    fns: LocationIntener<ItemLoc<ast::FnDef>, FunctionId>,
    structs: LocationIntener<ItemLoc<ast::StructDef>, StructId>,
    enums: LocationIntener<ItemLoc<ast::EnumDef>, EnumId>,
    consts: LocationIntener<ItemLoc<ast::ConstDef>, ConstId>,
    statics: LocationIntener<ItemLoc<ast::StaticDef>, StaticId>,
    traits: LocationIntener<ItemLoc<ast::TraitDef>, TraitId>,
    types: LocationIntener<ItemLoc<ast::TypeDef>, TypeId>,
}

impl HirInterner {
    pub fn len(&self) -> usize {
        self.macros.len()
            + self.fns.len()
            + self.structs.len()
            + self.enums.len()
            + self.consts.len()
            + self.statics.len()
            + self.traits.len()
            + self.types.len()
    }
}

/// hir makes heavy use of ids: integer (u32) handlers to various things. You
/// can think of id as a pointer (but without a lifetime) or a file descriptor
/// (but for hir objects).
///
/// This module defines a bunch of ids we are using. The most important ones are
/// probably `HirFileId` and `DefId`.

/// Input to the analyzer is a set of files, where each file is indentified by
/// `FileId` and contains source code. However, another source of source code in
/// Rust are macros: each macro can be thought of as producing a "temporary
/// file". To assign an id to such a file, we use the id of the macro call that
/// produced the file. So, a `HirFileId` is either a `FileId` (source code
/// written by user), or a `MacroCallId` (source code produced by macro).
///
/// What is a `MacroCallId`? Simplifying, it's a `HirFileId` of a file containin
/// the call plus the offset of the macro call in the file. Note that this is a
/// recursive definition! However, the size_of of `HirFileId` is finite
/// (because everything bottoms out at the real `FileId`) and small
/// (`MacroCallId` uses the location interner).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct HirFileId(HirFileIdRepr);

impl HirFileId {
    /// For macro-expansion files, returns the file original source file the
    /// expansionoriginated from.
    pub fn original_file(self, db: &impl HirDatabase) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(macro_call_id) => {
                let loc = macro_call_id.loc(db);
                loc.source_item_id.file_id.original_file(db)
            }
        }
    }

    pub(crate) fn as_original_file(self) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(_r) => panic!("macro generated file: {:?}", self),
        }
    }

    pub(crate) fn as_macro_call_id(self) -> Option<MacroCallId> {
        match self.0 {
            HirFileIdRepr::Macro(it) => Some(it),
            _ => None,
        }
    }

    pub(crate) fn hir_parse(db: &impl HirDatabase, file_id: HirFileId) -> TreeArc<SourceFile> {
        match file_id.0 {
            HirFileIdRepr::File(file_id) => db.parse(file_id),
            HirFileIdRepr::Macro(m) => {
                if let Some(exp) = db.expand_macro_invocation(m) {
                    return exp.file();
                }
                // returning an empty string looks fishy...
                SourceFile::parse("")
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum HirFileIdRepr {
    File(FileId),
    Macro(MacroCallId),
}

impl From<FileId> for HirFileId {
    fn from(file_id: FileId) -> HirFileId {
        HirFileId(HirFileIdRepr::File(file_id))
    }
}

impl From<MacroCallId> for HirFileId {
    fn from(macro_call_id: MacroCallId) -> HirFileId {
        HirFileId(HirFileIdRepr::Macro(macro_call_id))
    }
}

/// `MacroCallId` identifies a particular macro invocation, like
/// `println!("Hello, {}", world)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroCallId(RawId);
impl_arena_id!(MacroCallId);

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MacroCallLoc {
    pub(crate) module: Module,
    pub(crate) source_item_id: SourceItemId,
}

impl MacroCallId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> MacroCallLoc {
        db.as_ref().macros.id2loc(self)
    }
}

impl MacroCallLoc {
    #[allow(unused)]
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> MacroCallId {
        db.as_ref().macros.loc2id(&self)
    }
}

#[derive(Debug)]
pub struct ItemLoc<N: AstNode> {
    pub(crate) module: Module,
    raw: SourceItemId,
    _ty: PhantomData<N>,
}

impl<N: AstNode> PartialEq for ItemLoc<N> {
    fn eq(&self, other: &Self) -> bool {
        self.module == other.module && self.raw == other.raw
    }
}
impl<N: AstNode> Eq for ItemLoc<N> {}
impl<N: AstNode> Hash for ItemLoc<N> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.module.hash(hasher);
        self.raw.hash(hasher);
    }
}

impl<N: AstNode> Clone for ItemLoc<N> {
    fn clone(&self) -> ItemLoc<N> {
        ItemLoc {
            module: self.module,
            raw: self.raw,
            _ty: PhantomData,
        }
    }
}

#[derive(Clone, Copy)]
pub(crate) struct LocationCtx<DB> {
    db: DB,
    module: Module,
    file_id: HirFileId,
}

impl<'a, DB: HirDatabase> LocationCtx<&'a DB> {
    pub(crate) fn new(db: &'a DB, module: Module, file_id: HirFileId) -> LocationCtx<&'a DB> {
        LocationCtx {
            db,
            module,
            file_id,
        }
    }
    pub(crate) fn to_def<N, DEF>(self, ast: &N) -> DEF
    where
        N: AstNode,
        DEF: AstItemDef<N>,
    {
        DEF::from_ast(self, ast)
    }
}

pub(crate) trait AstItemDef<N: AstNode>: ArenaId + Clone {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<N>, Self>;
    fn from_ast(ctx: LocationCtx<&impl HirDatabase>, ast: &N) -> Self {
        let items = ctx.db.file_items(ctx.file_id);
        let raw = SourceItemId {
            file_id: ctx.file_id,
            item_id: Some(items.id_of(ctx.file_id, ast.syntax())),
        };
        let loc = ItemLoc {
            module: ctx.module,
            raw,
            _ty: PhantomData,
        };

        Self::interner(ctx.db.as_ref()).loc2id(&loc)
    }
    fn source(self, db: &impl HirDatabase) -> (HirFileId, TreeArc<N>) {
        let int = Self::interner(db.as_ref());
        let loc = int.id2loc(self);
        let syntax = db.file_item(loc.raw);
        let ast = N::cast(&syntax)
            .unwrap_or_else(|| panic!("invalid ItemLoc: {:?}", loc.raw))
            .to_owned();
        (loc.raw.file_id, ast)
    }
    fn module(self, db: &impl HirDatabase) -> Module {
        let int = Self::interner(db.as_ref());
        let loc = int.id2loc(self);
        loc.module
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct FunctionId(RawId);
impl_arena_id!(FunctionId);
impl AstItemDef<ast::FnDef> for FunctionId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::FnDef>, Self> {
        &interner.fns
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct StructId(RawId);
impl_arena_id!(StructId);
impl AstItemDef<ast::StructDef> for StructId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::StructDef>, Self> {
        &interner.structs
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct EnumId(RawId);
impl_arena_id!(EnumId);
impl AstItemDef<ast::EnumDef> for EnumId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::EnumDef>, Self> {
        &interner.enums
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ConstId(RawId);
impl_arena_id!(ConstId);
impl AstItemDef<ast::ConstDef> for ConstId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::ConstDef>, Self> {
        &interner.consts
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct StaticId(RawId);
impl_arena_id!(StaticId);
impl AstItemDef<ast::StaticDef> for StaticId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::StaticDef>, Self> {
        &interner.statics
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TraitId(RawId);
impl_arena_id!(TraitId);
impl AstItemDef<ast::TraitDef> for TraitId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::TraitDef>, Self> {
        &interner.traits
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TypeId(RawId);
impl_arena_id!(TypeId);
impl AstItemDef<ast::TypeDef> for TypeId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::TypeDef>, Self> {
        &interner.types
    }
}

/// Identifier of item within a specific file. This is stable over reparses, so
/// it's OK to use it as a salsa key/value.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SourceFileItemId(RawId);
impl_arena_id!(SourceFileItemId);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SourceItemId {
    pub(crate) file_id: HirFileId,
    /// None for the whole file.
    pub(crate) item_id: Option<SourceFileItemId>,
}

/// Maps items' `SyntaxNode`s to `SourceFileItemId`s and back.
#[derive(Debug, PartialEq, Eq)]
pub struct SourceFileItems {
    file_id: HirFileId,
    arena: Arena<SourceFileItemId, TreeArc<SyntaxNode>>,
}

impl SourceFileItems {
    pub(crate) fn new(file_id: HirFileId, source_file: &SourceFile) -> SourceFileItems {
        let mut res = SourceFileItems {
            file_id,
            arena: Arena::default(),
        };
        res.init(source_file);
        res
    }

    fn init(&mut self, source_file: &SourceFile) {
        // By walking the tree in bread-first order we make sure that parents
        // get lower ids then children. That is, addding a new child does not
        // change parent's id. This means that, say, adding a new function to a
        // trait does not chage ids of top-level items, which helps caching.
        bfs(source_file.syntax(), |it| {
            if let Some(module_item) = ast::ModuleItem::cast(it) {
                self.alloc(module_item.syntax().to_owned());
            } else if let Some(macro_call) = ast::MacroCall::cast(it) {
                self.alloc(macro_call.syntax().to_owned());
            }
        })
    }

    fn alloc(&mut self, item: TreeArc<SyntaxNode>) -> SourceFileItemId {
        self.arena.alloc(item)
    }
    pub(crate) fn id_of(&self, file_id: HirFileId, item: &SyntaxNode) -> SourceFileItemId {
        assert_eq!(
            self.file_id, file_id,
            "SourceFileItems: wrong file, expected {:?}, got {:?}",
            self.file_id, file_id
        );
        self.id_of_unchecked(item)
    }
    pub(crate) fn id_of_unchecked(&self, item: &SyntaxNode) -> SourceFileItemId {
        if let Some((id, _)) = self.arena.iter().find(|(_id, i)| *i == item) {
            return id;
        }
        // This should not happen. Let's try to give a sensible diagnostics.
        if let Some((id, i)) = self.arena.iter().find(|(_id, i)| i.range() == item.range()) {
            // FIXME(#288): whyyy are we getting here?
            log::error!(
                "unequal syntax nodes with the same range:\n{:?}\n{:?}",
                item,
                i
            );
            return id;
        }
        panic!(
            "Can't find {:?} in SourceFileItems:\n{:?}",
            item,
            self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
        );
    }
    pub fn id_of_parse(&self) -> SourceFileItemId {
        let (id, _syntax) = self.arena.iter().next().unwrap();
        id
    }
}

impl std::ops::Index<SourceFileItemId> for SourceFileItems {
    type Output = SyntaxNode;
    fn index(&self, idx: SourceFileItemId) -> &SyntaxNode {
        &self.arena[idx]
    }
}

/// Walks the subtree in bfs order, calling `f` for each node.
fn bfs(node: &SyntaxNode, mut f: impl FnMut(&SyntaxNode)) {
    let mut curr_layer = vec![node];
    let mut next_layer = vec![];
    while !curr_layer.is_empty() {
        curr_layer.drain(..).for_each(|node| {
            next_layer.extend(node.children());
            f(node);
        });
        std::mem::swap(&mut curr_layer, &mut next_layer);
    }
}
<commit_msg>remove dead code<commit_after>use std::{
    marker::PhantomData,
    hash::{Hash, Hasher},
};

use ra_db::{LocationIntener, FileId};
use ra_syntax::{TreeArc, SyntaxNode, SourceFile, AstNode, ast};
use ra_arena::{Arena, RawId, ArenaId, impl_arena_id};

use crate::{
    HirDatabase,
    Module,
};

#[derive(Debug, Default)]
pub struct HirInterner {
    macros: LocationIntener<MacroCallLoc, MacroCallId>,
    fns: LocationIntener<ItemLoc<ast::FnDef>, FunctionId>,
    structs: LocationIntener<ItemLoc<ast::StructDef>, StructId>,
    enums: LocationIntener<ItemLoc<ast::EnumDef>, EnumId>,
    consts: LocationIntener<ItemLoc<ast::ConstDef>, ConstId>,
    statics: LocationIntener<ItemLoc<ast::StaticDef>, StaticId>,
    traits: LocationIntener<ItemLoc<ast::TraitDef>, TraitId>,
    types: LocationIntener<ItemLoc<ast::TypeDef>, TypeId>,
}

impl HirInterner {
    pub fn len(&self) -> usize {
        self.macros.len()
            + self.fns.len()
            + self.structs.len()
            + self.enums.len()
            + self.consts.len()
            + self.statics.len()
            + self.traits.len()
            + self.types.len()
    }
}

/// hir makes heavy use of ids: integer (u32) handlers to various things. You
/// can think of id as a pointer (but without a lifetime) or a file descriptor
/// (but for hir objects).
///
/// This module defines a bunch of ids we are using. The most important ones are
/// probably `HirFileId` and `DefId`.

/// Input to the analyzer is a set of files, where each file is indentified by
/// `FileId` and contains source code. However, another source of source code in
/// Rust are macros: each macro can be thought of as producing a "temporary
/// file". To assign an id to such a file, we use the id of the macro call that
/// produced the file. So, a `HirFileId` is either a `FileId` (source code
/// written by user), or a `MacroCallId` (source code produced by macro).
///
/// What is a `MacroCallId`? Simplifying, it's a `HirFileId` of a file containin
/// the call plus the offset of the macro call in the file. Note that this is a
/// recursive definition! However, the size_of of `HirFileId` is finite
/// (because everything bottoms out at the real `FileId`) and small
/// (`MacroCallId` uses the location interner).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct HirFileId(HirFileIdRepr);

impl HirFileId {
    /// For macro-expansion files, returns the file original source file the
    /// expansionoriginated from.
    pub fn original_file(self, db: &impl HirDatabase) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(macro_call_id) => {
                let loc = macro_call_id.loc(db);
                loc.source_item_id.file_id.original_file(db)
            }
        }
    }

    pub(crate) fn as_original_file(self) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(_r) => panic!("macro generated file: {:?}", self),
        }
    }

    pub(crate) fn as_macro_call_id(self) -> Option<MacroCallId> {
        match self.0 {
            HirFileIdRepr::Macro(it) => Some(it),
            _ => None,
        }
    }

    pub(crate) fn hir_parse(db: &impl HirDatabase, file_id: HirFileId) -> TreeArc<SourceFile> {
        match file_id.0 {
            HirFileIdRepr::File(file_id) => db.parse(file_id),
            HirFileIdRepr::Macro(m) => {
                if let Some(exp) = db.expand_macro_invocation(m) {
                    return exp.file();
                }
                // returning an empty string looks fishy...
                SourceFile::parse("")
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum HirFileIdRepr {
    File(FileId),
    Macro(MacroCallId),
}

impl From<FileId> for HirFileId {
    fn from(file_id: FileId) -> HirFileId {
        HirFileId(HirFileIdRepr::File(file_id))
    }
}

impl From<MacroCallId> for HirFileId {
    fn from(macro_call_id: MacroCallId) -> HirFileId {
        HirFileId(HirFileIdRepr::Macro(macro_call_id))
    }
}

/// `MacroCallId` identifies a particular macro invocation, like
/// `println!("Hello, {}", world)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroCallId(RawId);
impl_arena_id!(MacroCallId);

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MacroCallLoc {
    pub(crate) module: Module,
    pub(crate) source_item_id: SourceItemId,
}

impl MacroCallId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> MacroCallLoc {
        db.as_ref().macros.id2loc(self)
    }
}

impl MacroCallLoc {
    #[allow(unused)]
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> MacroCallId {
        db.as_ref().macros.loc2id(&self)
    }
}

#[derive(Debug)]
pub struct ItemLoc<N: AstNode> {
    pub(crate) module: Module,
    raw: SourceItemId,
    _ty: PhantomData<N>,
}

impl<N: AstNode> PartialEq for ItemLoc<N> {
    fn eq(&self, other: &Self) -> bool {
        self.module == other.module && self.raw == other.raw
    }
}
impl<N: AstNode> Eq for ItemLoc<N> {}
impl<N: AstNode> Hash for ItemLoc<N> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.module.hash(hasher);
        self.raw.hash(hasher);
    }
}

impl<N: AstNode> Clone for ItemLoc<N> {
    fn clone(&self) -> ItemLoc<N> {
        ItemLoc {
            module: self.module,
            raw: self.raw,
            _ty: PhantomData,
        }
    }
}

#[derive(Clone, Copy)]
pub(crate) struct LocationCtx<DB> {
    db: DB,
    module: Module,
    file_id: HirFileId,
}

impl<'a, DB: HirDatabase> LocationCtx<&'a DB> {
    pub(crate) fn new(db: &'a DB, module: Module, file_id: HirFileId) -> LocationCtx<&'a DB> {
        LocationCtx {
            db,
            module,
            file_id,
        }
    }
    pub(crate) fn to_def<N, DEF>(self, ast: &N) -> DEF
    where
        N: AstNode,
        DEF: AstItemDef<N>,
    {
        DEF::from_ast(self, ast)
    }
}

pub(crate) trait AstItemDef<N: AstNode>: ArenaId + Clone {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<N>, Self>;
    fn from_ast(ctx: LocationCtx<&impl HirDatabase>, ast: &N) -> Self {
        let items = ctx.db.file_items(ctx.file_id);
        let raw = SourceItemId {
            file_id: ctx.file_id,
            item_id: Some(items.id_of(ctx.file_id, ast.syntax())),
        };
        let loc = ItemLoc {
            module: ctx.module,
            raw,
            _ty: PhantomData,
        };

        Self::interner(ctx.db.as_ref()).loc2id(&loc)
    }
    fn source(self, db: &impl HirDatabase) -> (HirFileId, TreeArc<N>) {
        let int = Self::interner(db.as_ref());
        let loc = int.id2loc(self);
        let syntax = db.file_item(loc.raw);
        let ast = N::cast(&syntax)
            .unwrap_or_else(|| panic!("invalid ItemLoc: {:?}", loc.raw))
            .to_owned();
        (loc.raw.file_id, ast)
    }
    fn module(self, db: &impl HirDatabase) -> Module {
        let int = Self::interner(db.as_ref());
        let loc = int.id2loc(self);
        loc.module
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct FunctionId(RawId);
impl_arena_id!(FunctionId);
impl AstItemDef<ast::FnDef> for FunctionId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::FnDef>, Self> {
        &interner.fns
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct StructId(RawId);
impl_arena_id!(StructId);
impl AstItemDef<ast::StructDef> for StructId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::StructDef>, Self> {
        &interner.structs
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct EnumId(RawId);
impl_arena_id!(EnumId);
impl AstItemDef<ast::EnumDef> for EnumId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::EnumDef>, Self> {
        &interner.enums
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ConstId(RawId);
impl_arena_id!(ConstId);
impl AstItemDef<ast::ConstDef> for ConstId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::ConstDef>, Self> {
        &interner.consts
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct StaticId(RawId);
impl_arena_id!(StaticId);
impl AstItemDef<ast::StaticDef> for StaticId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::StaticDef>, Self> {
        &interner.statics
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TraitId(RawId);
impl_arena_id!(TraitId);
impl AstItemDef<ast::TraitDef> for TraitId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::TraitDef>, Self> {
        &interner.traits
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TypeId(RawId);
impl_arena_id!(TypeId);
impl AstItemDef<ast::TypeDef> for TypeId {
    fn interner(interner: &HirInterner) -> &LocationIntener<ItemLoc<ast::TypeDef>, Self> {
        &interner.types
    }
}

/// Identifier of item within a specific file. This is stable over reparses, so
/// it's OK to use it as a salsa key/value.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SourceFileItemId(RawId);
impl_arena_id!(SourceFileItemId);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SourceItemId {
    pub(crate) file_id: HirFileId,
    /// None for the whole file.
    pub(crate) item_id: Option<SourceFileItemId>,
}

/// Maps items' `SyntaxNode`s to `SourceFileItemId`s and back.
#[derive(Debug, PartialEq, Eq)]
pub struct SourceFileItems {
    file_id: HirFileId,
    arena: Arena<SourceFileItemId, TreeArc<SyntaxNode>>,
}

impl SourceFileItems {
    pub(crate) fn new(file_id: HirFileId, source_file: &SourceFile) -> SourceFileItems {
        let mut res = SourceFileItems {
            file_id,
            arena: Arena::default(),
        };
        res.init(source_file);
        res
    }

    fn init(&mut self, source_file: &SourceFile) {
        // By walking the tree in bread-first order we make sure that parents
        // get lower ids then children. That is, addding a new child does not
        // change parent's id. This means that, say, adding a new function to a
        // trait does not chage ids of top-level items, which helps caching.
        bfs(source_file.syntax(), |it| {
            if let Some(module_item) = ast::ModuleItem::cast(it) {
                self.alloc(module_item.syntax().to_owned());
            } else if let Some(macro_call) = ast::MacroCall::cast(it) {
                self.alloc(macro_call.syntax().to_owned());
            }
        })
    }

    fn alloc(&mut self, item: TreeArc<SyntaxNode>) -> SourceFileItemId {
        self.arena.alloc(item)
    }
    pub(crate) fn id_of(&self, file_id: HirFileId, item: &SyntaxNode) -> SourceFileItemId {
        assert_eq!(
            self.file_id, file_id,
            "SourceFileItems: wrong file, expected {:?}, got {:?}",
            self.file_id, file_id
        );
        self.id_of_unchecked(item)
    }
    pub(crate) fn id_of_unchecked(&self, item: &SyntaxNode) -> SourceFileItemId {
        if let Some((id, _)) = self.arena.iter().find(|(_id, i)| *i == item) {
            return id;
        }
        // This should not happen. Let's try to give a sensible diagnostics.
        if let Some((id, i)) = self.arena.iter().find(|(_id, i)| i.range() == item.range()) {
            // FIXME(#288): whyyy are we getting here?
            log::error!(
                "unequal syntax nodes with the same range:\n{:?}\n{:?}",
                item,
                i
            );
            return id;
        }
        panic!(
            "Can't find {:?} in SourceFileItems:\n{:?}",
            item,
            self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
        );
    }
}

impl std::ops::Index<SourceFileItemId> for SourceFileItems {
    type Output = SyntaxNode;
    fn index(&self, idx: SourceFileItemId) -> &SyntaxNode {
        &self.arena[idx]
    }
}

/// Walks the subtree in bfs order, calling `f` for each node.
fn bfs(node: &SyntaxNode, mut f: impl FnMut(&SyntaxNode)) {
    let mut curr_layer = vec![node];
    let mut next_layer = vec![];
    while !curr_layer.is_empty() {
        curr_layer.drain(..).for_each(|node| {
            next_layer.extend(node.children());
            f(node);
        });
        std::mem::swap(&mut curr_layer, &mut next_layer);
    }
}
<|endoftext|>


<commit_before>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! component which does this.
//! need to get it into the memory in the first place somehow. VFS is the
//!
//! It also is responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! whihc are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod arena;
mod io;

use std::{
    fmt,
    mem,
    thread,
    cmp::Reverse,
    path::{Path, PathBuf},
    ffi::OsStr,
    sync::Arc,
    fs,
};

use rustc_hash::{FxHashMap, FxHashSet};
use relative_path::RelativePathBuf;
use crossbeam_channel::Receiver;
use walkdir::DirEntry;
use thread_worker::{WorkerHandle};

use crate::{
    arena::{ArenaId, Arena},
};

pub use crate::io::TaskResult as VfsTask;

/// `RootFilter` is a predicate that checks if a file can belong to a root. If
/// several filters match a file (nested dirs), the most nested one wins.
struct RootFilter {
    root: PathBuf,
    file_filter: fn(&Path) -> bool,
}

impl RootFilter {
    fn new(root: PathBuf) -> RootFilter {
        RootFilter {
            root,
            file_filter: has_rs_extension,
        }
    }
    /// Check if this root can contain `path`. NB: even if this returns
    /// true, the `path` might actually be conained in some nested root.
    fn can_contain(&self, path: &Path) -> Option<RelativePathBuf> {
        if !(self.file_filter)(path) {
            return None;
        }
        if !(path.starts_with(&self.root)) {
            return None;
        }
        let path = path.strip_prefix(&self.root).unwrap();
        let path = RelativePathBuf::from_path(path).unwrap();
        Some(path)
    }
}

fn has_rs_extension(p: &Path) -> bool {
    p.extension() == Some(OsStr::new("rs"))
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsRoot(pub u32);

impl ArenaId for VfsRoot {
    fn from_u32(idx: u32) -> VfsRoot {
        VfsRoot(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsFile(pub u32);

impl ArenaId for VfsFile {
    fn from_u32(idx: u32) -> VfsFile {
        VfsFile(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arena<VfsRoot, RootFilter>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: FxHashMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: io::Worker,
    worker_handle: WorkerHandle,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("Vfs { ... }")
    }
}

impl Vfs {
    pub fn new(mut roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let (worker, worker_handle) = io::start();

        let mut res = Vfs {
            roots: Arena::default(),
            files: Arena::default(),
            root2files: FxHashMap::default(),
            worker,
            worker_handle,
            pending_changes: Vec::new(),
        };

        // A hack to make nesting work.
        roots.sort_by_key(|it| Reverse(it.as_os_str().len()));
        for (i, path) in roots.iter().enumerate() {
            let root = res.roots.alloc(RootFilter::new(path.clone()));
            res.root2files.insert(root, Default::default());
            let nested = roots[..i]
                .iter()
                .filter(|it| it.starts_with(path))
                .map(|it| it.clone())
                .collect::<Vec<_>>();
            let filter = move |entry: &DirEntry| {
                if entry.file_type().is_file() {
                    has_rs_extension(entry.path())
                } else {
                    nested.iter().all(|it| it != entry.path())
                }
            };
            let task = io::Task {
                root,
                path: path.clone(),
                filter: Box::new(filter),
            };
            res.worker.inp.send(task);
        }
        let roots = res.roots.iter().map(|(id, _)| id).collect();
        (res, roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots[root].root.clone()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots[self.files[file].root].root;
        rel_path.to_path(root_path)
    }

    pub fn file_for_path(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.add_file(root, rel_path.clone(), Arc::clone(&text));
                let change = VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path: rel_path,
                };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        &self.worker.out
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        let mut files = Vec::new();
        // While we were scanning the root in the backgound, a file might have
        // been open in the editor, so we need to account for that.
        let exising = self.root2files[&task.root]
            .iter()
            .map(|&file| (self.files[file].path.clone(), file))
            .collect::<FxHashMap<_, _>>();
        for (path, text) in task.files {
            if let Some(&file) = exising.get(&path) {
                let text = Arc::clone(&self.files[file].text);
                files.push((file, path, text));
                continue;
            }
            let text = Arc::new(text);
            let file = self.add_file(task.root, path.clone(), Arc::clone(&text));
            files.push((file, path, text));
        }

        let change = VfsChange::AddRoot {
            root: task.root,
            files,
        };
        self.pending_changes.push(change);
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let text = Arc::new(text);
            let change = if let Some(file) = file {
                res = Some(file);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                let file = self.add_file(root, path.clone(), Arc::clone(&text));
                res = Some(file);
                VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path,
                }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            let text = Arc::new(new_text);
            self.change_file(file, Arc::clone(&text));
            let change = VfsChange::ChangeFile { file, text };
            self.pending_changes.push(change);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let file = file.expect("can't remove a file which wasn't added");
            res = Some(file);
            let full_path = path.to_path(&self.roots[root].root);
            let change = if let Ok(text) = fs::read_to_string(&full_path) {
                let text = Arc::new(text);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                self.remove_file(file);
                VfsChange::RemoveFile { root, file, path }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    /// Sutdown the VFS and terminate the background watching thread.
    pub fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }

    fn add_file(&mut self, root: VfsRoot, path: RelativePathBuf, text: Arc<String>) -> VfsFile {
        let data = VfsFileData { root, path, text };
        let file = self.files.alloc(data);
        self.root2files.get_mut(&root).unwrap().insert(file);
        file
    }

    fn change_file(&mut self, file: VfsFile, new_text: Arc<String>) {
        self.files[file].text = new_text;
    }

    fn remove_file(&mut self, file: VfsFile) {
        //FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(&root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self
            .roots
            .iter()
            .find_map(|(root, data)| data.can_contain(path).map(|it| (root, it)))?;
        let file = self.root2files[&root]
            .iter()
            .map(|&it| it)
            .find(|&file| self.files[file].path == path);
        Some((root, path, file))
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot {
        root: VfsRoot,
        files: Vec<(VfsFile, RelativePathBuf, Arc<String>)>,
    },
    AddFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
        text: Arc<String>,
    },
    RemoveFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
    },
    ChangeFile {
        file: VfsFile,
        text: Arc<String>,
    },
}
<commit_msg>Fix a switched line in a comment<commit_after>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! need to get it into the memory in the first place somehow. VFS is the
//! component which does this.
//!
//! It also is responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! whihc are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod arena;
mod io;

use std::{
    fmt,
    mem,
    thread,
    cmp::Reverse,
    path::{Path, PathBuf},
    ffi::OsStr,
    sync::Arc,
    fs,
};

use rustc_hash::{FxHashMap, FxHashSet};
use relative_path::RelativePathBuf;
use crossbeam_channel::Receiver;
use walkdir::DirEntry;
use thread_worker::{WorkerHandle};

use crate::{
    arena::{ArenaId, Arena},
};

pub use crate::io::TaskResult as VfsTask;

/// `RootFilter` is a predicate that checks if a file can belong to a root. If
/// several filters match a file (nested dirs), the most nested one wins.
struct RootFilter {
    root: PathBuf,
    file_filter: fn(&Path) -> bool,
}

impl RootFilter {
    fn new(root: PathBuf) -> RootFilter {
        RootFilter {
            root,
            file_filter: has_rs_extension,
        }
    }
    /// Check if this root can contain `path`. NB: even if this returns
    /// true, the `path` might actually be conained in some nested root.
    fn can_contain(&self, path: &Path) -> Option<RelativePathBuf> {
        if !(self.file_filter)(path) {
            return None;
        }
        if !(path.starts_with(&self.root)) {
            return None;
        }
        let path = path.strip_prefix(&self.root).unwrap();
        let path = RelativePathBuf::from_path(path).unwrap();
        Some(path)
    }
}

fn has_rs_extension(p: &Path) -> bool {
    p.extension() == Some(OsStr::new("rs"))
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsRoot(pub u32);

impl ArenaId for VfsRoot {
    fn from_u32(idx: u32) -> VfsRoot {
        VfsRoot(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsFile(pub u32);

impl ArenaId for VfsFile {
    fn from_u32(idx: u32) -> VfsFile {
        VfsFile(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arena<VfsRoot, RootFilter>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: FxHashMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: io::Worker,
    worker_handle: WorkerHandle,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("Vfs { ... }")
    }
}

impl Vfs {
    pub fn new(mut roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let (worker, worker_handle) = io::start();

        let mut res = Vfs {
            roots: Arena::default(),
            files: Arena::default(),
            root2files: FxHashMap::default(),
            worker,
            worker_handle,
            pending_changes: Vec::new(),
        };

        // A hack to make nesting work.
        roots.sort_by_key(|it| Reverse(it.as_os_str().len()));
        for (i, path) in roots.iter().enumerate() {
            let root = res.roots.alloc(RootFilter::new(path.clone()));
            res.root2files.insert(root, Default::default());
            let nested = roots[..i]
                .iter()
                .filter(|it| it.starts_with(path))
                .map(|it| it.clone())
                .collect::<Vec<_>>();
            let filter = move |entry: &DirEntry| {
                if entry.file_type().is_file() {
                    has_rs_extension(entry.path())
                } else {
                    nested.iter().all(|it| it != entry.path())
                }
            };
            let task = io::Task {
                root,
                path: path.clone(),
                filter: Box::new(filter),
            };
            res.worker.inp.send(task);
        }
        let roots = res.roots.iter().map(|(id, _)| id).collect();
        (res, roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots[root].root.clone()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots[self.files[file].root].root;
        rel_path.to_path(root_path)
    }

    pub fn file_for_path(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.add_file(root, rel_path.clone(), Arc::clone(&text));
                let change = VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path: rel_path,
                };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        &self.worker.out
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        let mut files = Vec::new();
        // While we were scanning the root in the backgound, a file might have
        // been open in the editor, so we need to account for that.
        let exising = self.root2files[&task.root]
            .iter()
            .map(|&file| (self.files[file].path.clone(), file))
            .collect::<FxHashMap<_, _>>();
        for (path, text) in task.files {
            if let Some(&file) = exising.get(&path) {
                let text = Arc::clone(&self.files[file].text);
                files.push((file, path, text));
                continue;
            }
            let text = Arc::new(text);
            let file = self.add_file(task.root, path.clone(), Arc::clone(&text));
            files.push((file, path, text));
        }

        let change = VfsChange::AddRoot {
            root: task.root,
            files,
        };
        self.pending_changes.push(change);
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let text = Arc::new(text);
            let change = if let Some(file) = file {
                res = Some(file);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                let file = self.add_file(root, path.clone(), Arc::clone(&text));
                res = Some(file);
                VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path,
                }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            let text = Arc::new(new_text);
            self.change_file(file, Arc::clone(&text));
            let change = VfsChange::ChangeFile { file, text };
            self.pending_changes.push(change);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let file = file.expect("can't remove a file which wasn't added");
            res = Some(file);
            let full_path = path.to_path(&self.roots[root].root);
            let change = if let Ok(text) = fs::read_to_string(&full_path) {
                let text = Arc::new(text);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                self.remove_file(file);
                VfsChange::RemoveFile { root, file, path }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    /// Sutdown the VFS and terminate the background watching thread.
    pub fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }

    fn add_file(&mut self, root: VfsRoot, path: RelativePathBuf, text: Arc<String>) -> VfsFile {
        let data = VfsFileData { root, path, text };
        let file = self.files.alloc(data);
        self.root2files.get_mut(&root).unwrap().insert(file);
        file
    }

    fn change_file(&mut self, file: VfsFile, new_text: Arc<String>) {
        self.files[file].text = new_text;
    }

    fn remove_file(&mut self, file: VfsFile) {
        //FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(&root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self
            .roots
            .iter()
            .find_map(|(root, data)| data.can_contain(path).map(|it| (root, it)))?;
        let file = self.root2files[&root]
            .iter()
            .map(|&it| it)
            .find(|&file| self.files[file].path == path);
        Some((root, path, file))
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot {
        root: VfsRoot,
        files: Vec<(VfsFile, RelativePathBuf, Arc<String>)>,
    },
    AddFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
        text: Arc<String>,
    },
    RemoveFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
    },
    ChangeFile {
        file: VfsFile,
        text: Arc<String>,
    },
}
<|endoftext|>