<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::iter::range_step;
use stb_image::image as stb_image;
use png;
use util::vec::byte_swap;

// FIXME: Images must not be copied every frame. Instead we should atomically
// reference count them.
pub type Image = png::Image;

static TEST_IMAGE: &'static [u8] = include_bytes!("test.jpeg");

pub fn test_image_bin() -> Vec<u8> {
    TEST_IMAGE.iter().map(|&x| x).collect()
}

// TODO(pcwalton): Speed up with SIMD, or better yet, find some way to not do this.
fn byte_swap_and_premultiply(data: &mut [u8]) {
    let length = data.len();
    for i in range_step(0, length, 4) {
        let r = data[i + 2];
        let g = data[i + 1];
        let b = data[i + 0];
        let a = data[i + 3];
        data[i + 0] = ((r as u32) * (a as u32) / 255) as u8;
        data[i + 1] = ((g as u32) * (a as u32) / 255) as u8;
        data[i + 2] = ((b as u32) * (a as u32) / 255) as u8;
    }
}

pub fn load_from_memory(buffer: &[u8]) -> Option<Image> {
    if buffer.len() == 0 {
        return None;
    }

    if png::is_png(buffer) {
        match png::load_png_from_memory(buffer) {
            Ok(mut png_image) => {
                match png_image.pixels {
                    png::PixelsByColorType::RGB8(ref mut data) => byte_swap(data.as_mut_slice()),
                    png::PixelsByColorType::RGBA8(ref mut data) => {
                        byte_swap_and_premultiply(data.as_mut_slice())
                    }
                    _ => {}
                }
                Some(png_image)
            }
            Err(_err) => None,
        }
    } else {
        // For non-png images, we use stb_image
        // Can't remember why we do this. Maybe it's what cairo wants
        static FORCE_DEPTH: uint = 4;

        match stb_image::load_from_memory_with_depth(buffer, FORCE_DEPTH, true) {
            stb_image::LoadResult::ImageU8(mut image) => {
                assert!(image.depth == 4);
                byte_swap(image.data.as_mut_slice());
                Some(png::Image {
                    width: image.width as u32,
                    height: image.height as u32,
                    pixels: png::PixelsByColorType::RGBA8(image.data)
                })
            }
            stb_image::LoadResult::ImageF32(_image) => {
                error!("HDR images not implemented");
                None
            }
            stb_image::LoadResult::Error(e) => {
                error!("stb_image failed: {}", e);
                None
            }
        }
    }
}
<commit_msg>Fix #5176 by premultiplying the alpha channel to the color channels<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::iter::range_step;
use stb_image::image as stb_image;
use png;
use util::vec::byte_swap;

// FIXME: Images must not be copied every frame. Instead we should atomically
// reference count them.
pub type Image = png::Image;

static TEST_IMAGE: &'static [u8] = include_bytes!("test.jpeg");

pub fn test_image_bin() -> Vec<u8> {
    TEST_IMAGE.iter().map(|&x| x).collect()
}

// TODO(pcwalton): Speed up with SIMD, or better yet, find some way to not do this.
fn byte_swap_and_premultiply(data: &mut [u8]) {
    let length = data.len();
    for i in range_step(0, length, 4) {
        let r = data[i + 2];
        let g = data[i + 1];
        let b = data[i + 0];
        let a = data[i + 3];
        data[i + 0] = ((r as u32) * (a as u32) / 255) as u8;
        data[i + 1] = ((g as u32) * (a as u32) / 255) as u8;
        data[i + 2] = ((b as u32) * (a as u32) / 255) as u8;
    }
}

pub fn load_from_memory(buffer: &[u8]) -> Option<Image> {
    if buffer.len() == 0 {
        return None;
    }

    if png::is_png(buffer) {
        match png::load_png_from_memory(buffer) {
            Ok(mut png_image) => {
                match png_image.pixels {
                    png::PixelsByColorType::RGB8(ref mut data) => byte_swap(data.as_mut_slice()),
                    png::PixelsByColorType::RGBA8(ref mut data) => {
                        byte_swap_and_premultiply(data.as_mut_slice())
                    }
                    _ => {}
                }
                Some(png_image)
            }
            Err(_err) => None,
        }
    } else {
        // For non-png images, we use stb_image
        // Can't remember why we do this. Maybe it's what cairo wants
        static FORCE_DEPTH: uint = 4;

        match stb_image::load_from_memory_with_depth(buffer, FORCE_DEPTH, true) {
            stb_image::LoadResult::ImageU8(mut image) => {
                assert!(image.depth == 4);
                // handle gif separately because the alpha-channel has to be premultiplied
                if is_gif(buffer) {
                    byte_swap_and_premultiply(image.data.as_mut_slice());
                } else {
                    byte_swap(image.data.as_mut_slice());
                }
                Some(png::Image {
                    width: image.width as u32,
                    height: image.height as u32,
                    pixels: png::PixelsByColorType::RGBA8(image.data)
                })
            }
            stb_image::LoadResult::ImageF32(_image) => {
                error!("HDR images not implemented");
                None
            }
            stb_image::LoadResult::Error(e) => {
                error!("stb_image failed: {}", e);
                None
            }
        }
    }
}

fn is_gif(buffer: &[u8]) -> bool {
    match buffer {
        [b'G',b'I',b'F',b'8', n, b'a', ..] if n == b'7' || n == b'9' => true,
        _ => false
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::URLBinding::{self, URLMethods};
use dom::bindings::error::{Error, ErrorResult, Fallible};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{JS, MutNullableHeap, Root};
use dom::bindings::reflector::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::str::USVString;
use dom::urlhelper::UrlHelper;
use dom::urlsearchparams::URLSearchParams;
use std::borrow::ToOwned;
use std::default::Default;
use url::{Host, ParseResult, Url, UrlParser};
use util::str::DOMString;

// https://url.spec.whatwg.org/#url
#[dom_struct]
pub struct URL {
    reflector_: Reflector,

    // https://url.spec.whatwg.org/#concept-urlutils-url
    url: DOMRefCell<Url>,

    // https://url.spec.whatwg.org/#concept-urlutils-get-the-base
    base: Option<Url>,

    // https://url.spec.whatwg.org/#dom-url-searchparams
    search_params: MutNullableHeap<JS<URLSearchParams>>,
}

impl URL {
    fn new_inherited(url: Url, base: Option<Url>) -> URL {
        URL {
            reflector_: Reflector::new(),
            url: DOMRefCell::new(url),
            base: base,
            search_params: Default::default(),
        }
    }

    pub fn new(global: GlobalRef, url: Url, base: Option<Url>) -> Root<URL> {
        reflect_dom_object(box URL::new_inherited(url, base),
                           global, URLBinding::Wrap)
    }

    pub fn set_query(&self, query: String) {
        self.url.borrow_mut().query = Some(query);
    }
}

impl URL {
    // https://url.spec.whatwg.org/#constructors
    pub fn Constructor(global: GlobalRef, url: USVString,
                       base: Option<USVString>)
                       -> Fallible<Root<URL>> {
        let parsed_base = match base {
            None => {
                // Step 1.
                None
            },
            Some(base) =>
                // Step 2.1.
                match Url::parse(&base.0) {
                    Ok(base) => Some(base),
                    Err(error) => {
                        // Step 2.2.
                        return Err(Error::Type(format!("could not parse base: {}", error)));
                    }
                }
        };
        // Step 3.
        let parsed_url = match parse_with_base(url, parsed_base.as_ref()) {
            Ok(url) => url,
            Err(error) => {
                // Step 4.
                return Err(Error::Type(format!("could not parse URL: {}", error)));
            }
        };
        // Step 5: Skip (see step 8 below).
        // Steps 6-7.
        let result = URL::new(global, parsed_url, parsed_base);
        // Step 8: Instead of construcing a new `URLSearchParams` object here, construct it
        //         on-demand inside `URL::SearchParams`.
        // Step 9.
        Ok(result)
    }

    // https://url.spec.whatwg.org/#dom-url-domaintoasciidomain
    pub fn DomainToASCII(_: GlobalRef, origin: USVString) -> USVString {
        // Step 1.
        let ascii_domain = Host::parse(&origin.0);
        if let Ok(Host::Domain(string)) = ascii_domain {
            // Step 3.
            USVString(string.to_owned())
        } else {
            // Step 2.
            USVString("".to_owned())
        }
    }
}

impl URLMethods for URL {
    // https://url.spec.whatwg.org/#dom-url-hash
    fn Hash(&self) -> USVString {
        UrlHelper::Hash(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-hash
    fn SetHash(&self, value: USVString) {
        UrlHelper::SetHash(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-host
    fn Host(&self) -> USVString {
        UrlHelper::Host(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-host
    fn SetHost(&self, value: USVString) {
        UrlHelper::SetHost(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-hostname
    fn Hostname(&self) -> USVString {
        UrlHelper::Hostname(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-hostname
    fn SetHostname(&self, value: USVString) {
        UrlHelper::SetHostname(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-href
    fn Href(&self) -> USVString {
        UrlHelper::Href(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-href
    fn SetHref(&self, value: USVString) -> ErrorResult {
        match parse_with_base(value, self.base.as_ref()) {
            Ok(url) => {
                *self.url.borrow_mut() = url;
                Ok(())
            },
            Err(error) => {
                Err(Error::Type(format!("could not parse URL: {}", error)))
            },
        }
    }

    // https://url.spec.whatwg.org/#dom-url-password
    fn Password(&self) -> USVString {
        UrlHelper::Password(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-password
    fn SetPassword(&self, value: USVString) {
        UrlHelper::SetPassword(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-pathname
    fn Pathname(&self) -> USVString {
        UrlHelper::Pathname(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-pathname
    fn SetPathname(&self, value: USVString) {
        UrlHelper::SetPathname(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-port
    fn Port(&self) -> USVString {
        UrlHelper::Port(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-port
    fn SetPort(&self, value: USVString) {
        UrlHelper::SetPort(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-protocol
    fn Protocol(&self) -> USVString {
        UrlHelper::Protocol(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-protocol
    fn SetProtocol(&self, value: USVString) {
        UrlHelper::SetProtocol(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-origin
    fn Origin(&self) -> USVString {
        UrlHelper::Origin(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-search
    fn Search(&self) -> USVString {
        UrlHelper::Search(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-search
    fn SetSearch(&self, value: USVString) {
        UrlHelper::SetSearch(&mut self.url.borrow_mut(), value);
        if let Some(search_params) = self.search_params.get() {
            search_params.set_list(self.url.borrow().query_pairs().unwrap_or_else(|| vec![]));
        }
    }

    // https://url.spec.whatwg.org/#dom-url-searchparams
    fn SearchParams(&self) -> Root<URLSearchParams> {
        self.search_params.or_init(|| URLSearchParams::new(self.global().r(), Some(self)))
    }

    // https://url.spec.whatwg.org/#dom-url-href
    fn Stringifier(&self) -> DOMString {
        DOMString::from(self.Href().0)
    }

    // https://url.spec.whatwg.org/#dom-url-username
    fn Username(&self) -> USVString {
        UrlHelper::Username(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-username
    fn SetUsername(&self, value: USVString) {
        UrlHelper::SetUsername(&mut self.url.borrow_mut(), value);
    }
}

fn parse_with_base(input: USVString, base: Option<&Url>) -> ParseResult<Url> {
    let mut parser = UrlParser::new();
    if let Some(base) = base {
        parser.base_url(base);
    }
    parser.parse(&input.0)
}
<commit_msg>Auto merge of #10430 - stjepang:remove-url-base, r=Ms2ger<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::URLBinding::{self, URLMethods};
use dom::bindings::error::{Error, ErrorResult, Fallible};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{JS, MutNullableHeap, Root};
use dom::bindings::reflector::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::str::USVString;
use dom::urlhelper::UrlHelper;
use dom::urlsearchparams::URLSearchParams;
use std::borrow::ToOwned;
use std::default::Default;
use url::{Host, Url, UrlParser};
use util::str::DOMString;

// https://url.spec.whatwg.org/#url
#[dom_struct]
pub struct URL {
    reflector_: Reflector,

    // https://url.spec.whatwg.org/#concept-url-url
    url: DOMRefCell<Url>,

    // https://url.spec.whatwg.org/#dom-url-searchparams
    search_params: MutNullableHeap<JS<URLSearchParams>>,
}

impl URL {
    fn new_inherited(url: Url) -> URL {
        URL {
            reflector_: Reflector::new(),
            url: DOMRefCell::new(url),
            search_params: Default::default(),
        }
    }

    pub fn new(global: GlobalRef, url: Url) -> Root<URL> {
        reflect_dom_object(box URL::new_inherited(url),
                           global, URLBinding::Wrap)
    }

    pub fn set_query(&self, query: String) {
        self.url.borrow_mut().query = Some(query);
    }
}

impl URL {
    // https://url.spec.whatwg.org/#constructors
    pub fn Constructor(global: GlobalRef, url: USVString,
                       base: Option<USVString>)
                       -> Fallible<Root<URL>> {
        let parsed_base = match base {
            None => {
                // Step 1.
                None
            },
            Some(base) =>
                // Step 2.1.
                match Url::parse(&base.0) {
                    Ok(base) => Some(base),
                    Err(error) => {
                        // Step 2.2.
                        return Err(Error::Type(format!("could not parse base: {}", error)));
                    }
                }
        };
        // Step 3.
        let parsed_url = {
            let mut parser = UrlParser::new();
            if let Some(parsed_base) = parsed_base.as_ref() {
                parser.base_url(parsed_base);
            }
            match parser.parse(&url.0) {
                Ok(url) => url,
                Err(error) => {
                    // Step 4.
                    return Err(Error::Type(format!("could not parse URL: {}", error)));
                }
            }
        };
        // Step 5: Skip (see step 8 below).
        // Steps 6-7.
        let result = URL::new(global, parsed_url);
        // Step 8: Instead of construcing a new `URLSearchParams` object here, construct it
        //         on-demand inside `URL::SearchParams`.
        // Step 9.
        Ok(result)
    }

    // https://url.spec.whatwg.org/#dom-url-domaintoasciidomain
    pub fn DomainToASCII(_: GlobalRef, origin: USVString) -> USVString {
        // Step 1.
        let ascii_domain = Host::parse(&origin.0);
        if let Ok(Host::Domain(string)) = ascii_domain {
            // Step 3.
            USVString(string.to_owned())
        } else {
            // Step 2.
            USVString("".to_owned())
        }
    }
}

impl URLMethods for URL {
    // https://url.spec.whatwg.org/#dom-url-hash
    fn Hash(&self) -> USVString {
        UrlHelper::Hash(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-hash
    fn SetHash(&self, value: USVString) {
        UrlHelper::SetHash(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-host
    fn Host(&self) -> USVString {
        UrlHelper::Host(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-host
    fn SetHost(&self, value: USVString) {
        UrlHelper::SetHost(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-hostname
    fn Hostname(&self) -> USVString {
        UrlHelper::Hostname(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-hostname
    fn SetHostname(&self, value: USVString) {
        UrlHelper::SetHostname(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-href
    fn Href(&self) -> USVString {
        UrlHelper::Href(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-href
    fn SetHref(&self, value: USVString) -> ErrorResult {
        match Url::parse(&value.0) {
            Ok(url) => {
                *self.url.borrow_mut() = url;
                Ok(())
            },
            Err(error) => {
                Err(Error::Type(format!("could not parse URL: {}", error)))
            },
        }
    }

    // https://url.spec.whatwg.org/#dom-url-password
    fn Password(&self) -> USVString {
        UrlHelper::Password(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-password
    fn SetPassword(&self, value: USVString) {
        UrlHelper::SetPassword(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-pathname
    fn Pathname(&self) -> USVString {
        UrlHelper::Pathname(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-pathname
    fn SetPathname(&self, value: USVString) {
        UrlHelper::SetPathname(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-port
    fn Port(&self) -> USVString {
        UrlHelper::Port(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-port
    fn SetPort(&self, value: USVString) {
        UrlHelper::SetPort(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-protocol
    fn Protocol(&self) -> USVString {
        UrlHelper::Protocol(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-protocol
    fn SetProtocol(&self, value: USVString) {
        UrlHelper::SetProtocol(&mut self.url.borrow_mut(), value);
    }

    // https://url.spec.whatwg.org/#dom-url-origin
    fn Origin(&self) -> USVString {
        UrlHelper::Origin(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-search
    fn Search(&self) -> USVString {
        UrlHelper::Search(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-search
    fn SetSearch(&self, value: USVString) {
        UrlHelper::SetSearch(&mut self.url.borrow_mut(), value);
        if let Some(search_params) = self.search_params.get() {
            search_params.set_list(self.url.borrow().query_pairs().unwrap_or_else(|| vec![]));
        }
    }

    // https://url.spec.whatwg.org/#dom-url-searchparams
    fn SearchParams(&self) -> Root<URLSearchParams> {
        self.search_params.or_init(|| URLSearchParams::new(self.global().r(), Some(self)))
    }

    // https://url.spec.whatwg.org/#dom-url-href
    fn Stringifier(&self) -> DOMString {
        DOMString::from(self.Href().0)
    }

    // https://url.spec.whatwg.org/#dom-url-username
    fn Username(&self) -> USVString {
        UrlHelper::Username(&self.url.borrow())
    }

    // https://url.spec.whatwg.org/#dom-url-username
    fn SetUsername(&self, value: USVString) {
        UrlHelper::SetUsername(&mut self.url.borrow_mut(), value);
    }
}
<|endoftext|>


<commit_before>#![feature(plugin_registrar, quote, rustc_private, custom_attribute)]

extern crate rustc_plugin;
extern crate syntax;

use rustc_plugin::registry::Registry;
use syntax::ast::{Attribute, Block, Expr, ExprKind, Ident, Item, ItemKind, Mac,
                  MetaItem};
use syntax::fold::{self, Folder};
use syntax::ptr::P;
use syntax::codemap::{DUMMY_SP, Span};
use syntax::ext::base::{Annotatable, ExtCtxt, SyntaxExtension};
use syntax::ext::build::AstBuilder;
use syntax::feature_gate::AttributeType;
use syntax::symbol::Symbol;
use syntax::util::small_vector::SmallVector;

pub fn insert_flame_guard(cx: &mut ExtCtxt, _span: Span, _mi: &MetaItem,
                          a: Annotatable) -> Annotatable {
    match a {
        Annotatable::Item(i) => Annotatable::Item(
            Flamer { cx: cx, ident: i.ident }.fold_item(i).expect_one("expected exactly one item")),
        Annotatable::TraitItem(i) => Annotatable::TraitItem(
            i.map(|i| Flamer { cx: cx, ident: i.ident }.fold_trait_item(i).expect_one("expected exactly one item"))),
        Annotatable::ImplItem(i) => Annotatable::ImplItem(
            i.map(|i| Flamer { cx: cx, ident: i.ident }.fold_impl_item(i).expect_one("expected exactly one item"))),
    }
}

struct Flamer<'a, 'cx: 'a> {
    ident: Ident,
    cx: &'a mut ExtCtxt<'cx>,
}

impl<'a, 'cx> Folder for Flamer<'a, 'cx> {
    fn fold_item(&mut self, item: P<Item>) -> SmallVector<P<Item>> {
        if let ItemKind::Mac(_) = item.node {
            let expanded = self.cx.expander().fold_item(item);
            expanded.into_iter()
                    .flat_map(|i| fold::noop_fold_item(i, self).into_iter())
                    .collect()
        } else {
            fold::noop_fold_item(item, self)
        }
    }

    fn fold_item_simple(&mut self, i: Item) -> Item {
        fn is_flame_annotation(attr: &Attribute) -> bool {
            attr.name().map_or(false, |name|
                    name == "flame" || name == "noflame")
        }
        // don't double-flame nested annotations
        if i.attrs.iter().any(is_flame_annotation) { return i; }
        if let ItemKind::Mac(_) = i.node {
            return i;
        } else {
            self.ident = i.ident; // update in case of nested items
            fold::noop_fold_item_simple(i, self)
        }
    }

    fn fold_block(&mut self, block: P<Block>) -> P<Block> {
        block.map(|block| {
            let name = self.cx.expr_str(DUMMY_SP, self.ident.name);
            quote_block!(self.cx, {
                let g = ::profiler::ProfilerSpan::new($name);
                let r = $block;
                drop(g);
                r
            }).unwrap()

        })
    }

    fn fold_expr(&mut self, expr: P<Expr>) -> P<Expr> {
        if let ExprKind::Mac(_) = expr.node {
            self.cx.expander().fold_expr(expr)
                              .map(|e| fold::noop_fold_expr(e, self))
        } else {
            expr
        }
    }

    fn fold_mac(&mut self, mac: Mac) -> Mac {
        mac
    }
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(Symbol::intern("flame"),
        SyntaxExtension::MultiModifier(Box::new(insert_flame_guard)));
    reg.register_attribute(String::from("noflame"), AttributeType::Whitelisted);
}
<commit_msg>allow unused imports<commit_after>#![allow(unused_imports)]
#![feature(plugin_registrar, quote, rustc_private, custom_attribute)]

extern crate rustc_plugin;
extern crate syntax;

use rustc_plugin::registry::Registry;
use syntax::ast::{Attribute, Block, Expr, ExprKind, Ident, Item, ItemKind, Mac,
                  MetaItem};
use syntax::fold::{self, Folder};
use syntax::ptr::P;
use syntax::codemap::{DUMMY_SP, Span};
use syntax::ext::base::{Annotatable, ExtCtxt, SyntaxExtension};
use syntax::ext::build::AstBuilder;
use syntax::feature_gate::AttributeType;
use syntax::symbol::Symbol;
use syntax::util::small_vector::SmallVector;

pub fn insert_flame_guard(cx: &mut ExtCtxt, _span: Span, _mi: &MetaItem,
                          a: Annotatable) -> Annotatable {
    match a {
        Annotatable::Item(i) => Annotatable::Item(
            Flamer { cx: cx, ident: i.ident }.fold_item(i).expect_one("expected exactly one item")),
        Annotatable::TraitItem(i) => Annotatable::TraitItem(
            i.map(|i| Flamer { cx: cx, ident: i.ident }.fold_trait_item(i).expect_one("expected exactly one item"))),
        Annotatable::ImplItem(i) => Annotatable::ImplItem(
            i.map(|i| Flamer { cx: cx, ident: i.ident }.fold_impl_item(i).expect_one("expected exactly one item"))),
    }
}

struct Flamer<'a, 'cx: 'a> {
    ident: Ident,
    cx: &'a mut ExtCtxt<'cx>,
}

impl<'a, 'cx> Folder for Flamer<'a, 'cx> {
    fn fold_item(&mut self, item: P<Item>) -> SmallVector<P<Item>> {
        if let ItemKind::Mac(_) = item.node {
            let expanded = self.cx.expander().fold_item(item);
            expanded.into_iter()
                    .flat_map(|i| fold::noop_fold_item(i, self).into_iter())
                    .collect()
        } else {
            fold::noop_fold_item(item, self)
        }
    }

    fn fold_item_simple(&mut self, i: Item) -> Item {
        fn is_flame_annotation(attr: &Attribute) -> bool {
            attr.name().map_or(false, |name|
                    name == "flame" || name == "noflame")
        }
        // don't double-flame nested annotations
        if i.attrs.iter().any(is_flame_annotation) { return i; }
        if let ItemKind::Mac(_) = i.node {
            return i;
        } else {
            self.ident = i.ident; // update in case of nested items
            fold::noop_fold_item_simple(i, self)
        }
    }

    fn fold_block(&mut self, block: P<Block>) -> P<Block> {
        block.map(|block| {
            let name = self.cx.expr_str(DUMMY_SP, self.ident.name);
            quote_block!(self.cx, {
                let g = ::profiler::ProfilerSpan::new($name);
                let r = $block;
                drop(g);
                r
            }).unwrap()

        })
    }

    fn fold_expr(&mut self, expr: P<Expr>) -> P<Expr> {
        if let ExprKind::Mac(_) = expr.node {
            self.cx.expander().fold_expr(expr)
                              .map(|e| fold::noop_fold_expr(e, self))
        } else {
            expr
        }
    }

    fn fold_mac(&mut self, mac: Mac) -> Mac {
        mac
    }
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(Symbol::intern("flame"),
        SyntaxExtension::MultiModifier(Box::new(insert_flame_guard)));
    reg.register_attribute(String::from("noflame"), AttributeType::Whitelisted);
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! The type_coercion optimizer rule ensures that all operators are operating on
//! compatible types by adding explicit cast operations to expressions. For example,
//! the operation `c_float + c_int` would be rewritten as `c_float + CAST(c_int AS
//! float)`. This keeps the runtime query execution code much simpler.

use std::collections::HashMap;

use arrow::datatypes::Schema;

use crate::error::{ExecutionError, Result};
use crate::execution::physical_plan::udf::ScalarFunction;
use crate::logicalplan::LogicalPlan;
use crate::logicalplan::{Expr, LogicalPlanBuilder};
use crate::optimizer::optimizer::OptimizerRule;
use crate::optimizer::utils;

/// Implementation of type coercion optimizer rule
pub struct TypeCoercionRule<'a> {
    scalar_functions: &'a HashMap<String, Box<ScalarFunction>>,
}

impl<'a> TypeCoercionRule<'a> {
    /// Create a new type coercion optimizer rule using meta-data about registered
    /// scalar functions
    pub fn new(scalar_functions: &'a HashMap<String, Box<ScalarFunction>>) -> Self {
        Self { scalar_functions }
    }

    /// Rewrite an expression list to include explicit CAST operations when required
    fn rewrite_expr_list(&self, expr: &[Expr], schema: &Schema) -> Result<Vec<Expr>> {
        Ok(expr
            .iter()
            .map(|e| self.rewrite_expr(e, schema))
            .collect::<Result<Vec<_>>>()?)
    }

    /// Rewrite an expression to include explicit CAST operations when required
    fn rewrite_expr(&self, expr: &Expr, schema: &Schema) -> Result<Expr> {
        match expr {
            Expr::BinaryExpr { left, op, right } => {
                let left = self.rewrite_expr(left, schema)?;
                let right = self.rewrite_expr(right, schema)?;
                let left_type = left.get_type(schema)?;
                let right_type = right.get_type(schema)?;
                if left_type == right_type {
                    Ok(Expr::BinaryExpr {
                        left: Box::new(left),
                        op: op.clone(),
                        right: Box::new(right),
                    })
                } else {
                    let super_type = utils::get_supertype(&left_type, &right_type)?;
                    Ok(Expr::BinaryExpr {
                        left: Box::new(left.cast_to(&super_type, schema)?),
                        op: op.clone(),
                        right: Box::new(right.cast_to(&super_type, schema)?),
                    })
                }
            }
            Expr::IsNull(e) => Ok(Expr::IsNull(Box::new(self.rewrite_expr(e, schema)?))),
            Expr::IsNotNull(e) => {
                Ok(Expr::IsNotNull(Box::new(self.rewrite_expr(e, schema)?)))
            }
            Expr::ScalarFunction {
                name,
                args,
                return_type,
            } => {
                // cast the inputs of scalar functions to the appropriate type where possible
                match self.scalar_functions.get(name) {
                    Some(func_meta) => {
                        let mut func_args = Vec::with_capacity(args.len());
                        for i in 0..args.len() {
                            let field = &func_meta.args[i];
                            let expr = self.rewrite_expr(&args[i], schema)?;
                            let actual_type = expr.get_type(schema)?;
                            let required_type = field.data_type();
                            if &actual_type == required_type {
                                func_args.push(expr)
                            } else {
                                let super_type =
                                    utils::get_supertype(&actual_type, required_type)?;
                                func_args.push(expr.cast_to(&super_type, schema)?);
                            }
                        }

                        Ok(Expr::ScalarFunction {
                            name: name.clone(),
                            args: func_args,
                            return_type: return_type.clone(),
                        })
                    }
                    _ => Err(ExecutionError::General(format!(
                        "Invalid scalar function {}",
                        name
                    ))),
                }
            }
            Expr::AggregateFunction {
                name,
                args,
                return_type,
            } => Ok(Expr::AggregateFunction {
                name: name.clone(),
                args: args
                    .iter()
                    .map(|a| self.rewrite_expr(a, schema))
                    .collect::<Result<Vec<_>>>()?,
                return_type: return_type.clone(),
            }),
            Expr::Cast { .. } => Ok(expr.clone()),
            Expr::Column(_) => Ok(expr.clone()),
            Expr::Alias(expr, alias) => Ok(Expr::Alias(
                Box::new(self.rewrite_expr(expr, schema)?),
                alias.to_owned(),
            )),
            Expr::Literal(_) => Ok(expr.clone()),
            Expr::Not(_) => Ok(expr.clone()),
            Expr::Sort { .. } => Ok(expr.clone()),
            Expr::Wildcard { .. } => Err(ExecutionError::General(
                "Wildcard expressions are not valid in a logical query plan".to_owned(),
            )),
            Expr::Nested(e) => self.rewrite_expr(e, schema),
        }
    }
}

impl<'a> OptimizerRule for TypeCoercionRule<'a> {
    fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
        match plan {
            LogicalPlan::Projection { expr, input, .. } => {
                LogicalPlanBuilder::from(&self.optimize(input)?)
                    .project(self.rewrite_expr_list(expr, input.schema())?)?
                    .build()
            }
            LogicalPlan::Selection { expr, input, .. } => {
                LogicalPlanBuilder::from(&self.optimize(input)?)
                    .filter(self.rewrite_expr(expr, input.schema())?)?
                    .build()
            }
            LogicalPlan::Aggregate {
                input,
                group_expr,
                aggr_expr,
                ..
            } => LogicalPlanBuilder::from(&self.optimize(input)?)
                .aggregate(
                    self.rewrite_expr_list(group_expr, input.schema())?,
                    self.rewrite_expr_list(aggr_expr, input.schema())?,
                )?
                .build(),
            LogicalPlan::TableScan { .. } => Ok(plan.clone()),
            LogicalPlan::InMemoryScan { .. } => Ok(plan.clone()),
            LogicalPlan::ParquetScan { .. } => Ok(plan.clone()),
            LogicalPlan::CsvScan { .. } => Ok(plan.clone()),
            LogicalPlan::EmptyRelation { .. } => Ok(plan.clone()),
            LogicalPlan::Limit { .. } => Ok(plan.clone()),
            LogicalPlan::Sort { .. } => Ok(plan.clone()),
            LogicalPlan::CreateExternalTable { .. } => Ok(plan.clone()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::execution::context::ExecutionContext;
    use crate::execution::physical_plan::csv::CsvReadOptions;
    use crate::logicalplan::{col, Operator};
    use crate::test::arrow_testdata_path;
    use arrow::datatypes::{DataType, Field, Schema};

    #[test]
    fn test_with_csv_plan() -> Result<()> {
        let testdata = arrow_testdata_path();
        let path = format!("{}/csv/aggregate_test_100.csv", testdata);

        let options = CsvReadOptions::new().schema_infer_max_records(100);
        let plan = LogicalPlanBuilder::scan_csv(&path, options, None)?
            .filter(col("c7").lt(&col("c12")))?
            .build()?;

        let scalar_functions = HashMap::new();
        let mut rule = TypeCoercionRule::new(&scalar_functions);
        let plan = rule.optimize(&plan)?;

        assert!(
            format!("{:?}", plan).starts_with("Selection: CAST(#c7 AS Float64) Lt #c12")
        );

        Ok(())
    }

    #[test]
    fn test_add_i32_i64() {
        binary_cast_test(
            DataType::Int32,
            DataType::Int64,
            "CAST(#c0 AS Int64) Plus #c1",
        );
        binary_cast_test(
            DataType::Int64,
            DataType::Int32,
            "#c0 Plus CAST(#c1 AS Int64)",
        );
    }

    #[test]
    fn test_add_f32_f64() {
        binary_cast_test(
            DataType::Float32,
            DataType::Float64,
            "CAST(#c0 AS Float64) Plus #c1",
        );
        binary_cast_test(
            DataType::Float64,
            DataType::Float32,
            "#c0 Plus CAST(#c1 AS Float64)",
        );
    }

    #[test]
    fn test_add_i32_f32() {
        binary_cast_test(
            DataType::Int32,
            DataType::Float32,
            "CAST(#c0 AS Float32) Plus #c1",
        );
        binary_cast_test(
            DataType::Float32,
            DataType::Int32,
            "#c0 Plus CAST(#c1 AS Float32)",
        );
    }

    #[test]
    fn test_add_u32_i64() {
        binary_cast_test(
            DataType::UInt32,
            DataType::Int64,
            "CAST(#c0 AS Int64) Plus #c1",
        );
        binary_cast_test(
            DataType::Int64,
            DataType::UInt32,
            "#c0 Plus CAST(#c1 AS Int64)",
        );
    }

    fn binary_cast_test(left_type: DataType, right_type: DataType, expected: &str) {
        let schema = Schema::new(vec![
            Field::new("c0", left_type, true),
            Field::new("c1", right_type, true),
        ]);

        let expr = Expr::BinaryExpr {
            left: Box::new(col("c0")),
            op: Operator::Plus,
            right: Box::new(col("c1")),
        };

        let ctx = ExecutionContext::new();
        let rule = TypeCoercionRule::new(ctx.scalar_functions());

        let expr2 = rule.rewrite_expr(&expr, &schema).unwrap();

        assert_eq!(expected, format!("{:?}", expr2));
    }
}
<commit_msg>ARROW-9714: [Rust] [DataFusion] Implement type coercion rule for limit and sort<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! The type_coercion optimizer rule ensures that all operators are operating on
//! compatible types by adding explicit cast operations to expressions. For example,
//! the operation `c_float + c_int` would be rewritten as `c_float + CAST(c_int AS
//! float)`. This keeps the runtime query execution code much simpler.

use std::collections::HashMap;

use arrow::datatypes::Schema;

use crate::error::{ExecutionError, Result};
use crate::execution::physical_plan::udf::ScalarFunction;
use crate::logicalplan::LogicalPlan;
use crate::logicalplan::{Expr, LogicalPlanBuilder};
use crate::optimizer::optimizer::OptimizerRule;
use crate::optimizer::utils;

/// Implementation of type coercion optimizer rule
pub struct TypeCoercionRule<'a> {
    scalar_functions: &'a HashMap<String, Box<ScalarFunction>>,
}

impl<'a> TypeCoercionRule<'a> {
    /// Create a new type coercion optimizer rule using meta-data about registered
    /// scalar functions
    pub fn new(scalar_functions: &'a HashMap<String, Box<ScalarFunction>>) -> Self {
        Self { scalar_functions }
    }

    /// Rewrite an expression list to include explicit CAST operations when required
    fn rewrite_expr_list(&self, expr: &[Expr], schema: &Schema) -> Result<Vec<Expr>> {
        Ok(expr
            .iter()
            .map(|e| self.rewrite_expr(e, schema))
            .collect::<Result<Vec<_>>>()?)
    }

    /// Rewrite an expression to include explicit CAST operations when required
    fn rewrite_expr(&self, expr: &Expr, schema: &Schema) -> Result<Expr> {
        match expr {
            Expr::BinaryExpr { left, op, right } => {
                let left = self.rewrite_expr(left, schema)?;
                let right = self.rewrite_expr(right, schema)?;
                let left_type = left.get_type(schema)?;
                let right_type = right.get_type(schema)?;
                if left_type == right_type {
                    Ok(Expr::BinaryExpr {
                        left: Box::new(left),
                        op: op.clone(),
                        right: Box::new(right),
                    })
                } else {
                    let super_type = utils::get_supertype(&left_type, &right_type)?;
                    Ok(Expr::BinaryExpr {
                        left: Box::new(left.cast_to(&super_type, schema)?),
                        op: op.clone(),
                        right: Box::new(right.cast_to(&super_type, schema)?),
                    })
                }
            }
            Expr::IsNull(e) => Ok(Expr::IsNull(Box::new(self.rewrite_expr(e, schema)?))),
            Expr::IsNotNull(e) => {
                Ok(Expr::IsNotNull(Box::new(self.rewrite_expr(e, schema)?)))
            }
            Expr::ScalarFunction {
                name,
                args,
                return_type,
            } => {
                // cast the inputs of scalar functions to the appropriate type where possible
                match self.scalar_functions.get(name) {
                    Some(func_meta) => {
                        let mut func_args = Vec::with_capacity(args.len());
                        for i in 0..args.len() {
                            let field = &func_meta.args[i];
                            let expr = self.rewrite_expr(&args[i], schema)?;
                            let actual_type = expr.get_type(schema)?;
                            let required_type = field.data_type();
                            if &actual_type == required_type {
                                func_args.push(expr)
                            } else {
                                let super_type =
                                    utils::get_supertype(&actual_type, required_type)?;
                                func_args.push(expr.cast_to(&super_type, schema)?);
                            }
                        }

                        Ok(Expr::ScalarFunction {
                            name: name.clone(),
                            args: func_args,
                            return_type: return_type.clone(),
                        })
                    }
                    _ => Err(ExecutionError::General(format!(
                        "Invalid scalar function {}",
                        name
                    ))),
                }
            }
            Expr::AggregateFunction {
                name,
                args,
                return_type,
            } => Ok(Expr::AggregateFunction {
                name: name.clone(),
                args: args
                    .iter()
                    .map(|a| self.rewrite_expr(a, schema))
                    .collect::<Result<Vec<_>>>()?,
                return_type: return_type.clone(),
            }),
            Expr::Cast { .. } => Ok(expr.clone()),
            Expr::Column(_) => Ok(expr.clone()),
            Expr::Alias(expr, alias) => Ok(Expr::Alias(
                Box::new(self.rewrite_expr(expr, schema)?),
                alias.to_owned(),
            )),
            Expr::Literal(_) => Ok(expr.clone()),
            Expr::Not(_) => Ok(expr.clone()),
            Expr::Sort { .. } => Ok(expr.clone()),
            Expr::Wildcard { .. } => Err(ExecutionError::General(
                "Wildcard expressions are not valid in a logical query plan".to_owned(),
            )),
            Expr::Nested(e) => self.rewrite_expr(e, schema),
        }
    }
}

impl<'a> OptimizerRule for TypeCoercionRule<'a> {
    fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
        match plan {
            LogicalPlan::Projection { expr, input, .. } => {
                LogicalPlanBuilder::from(&self.optimize(input)?)
                    .project(self.rewrite_expr_list(expr, input.schema())?)?
                    .build()
            }
            LogicalPlan::Selection { expr, input, .. } => {
                LogicalPlanBuilder::from(&self.optimize(input)?)
                    .filter(self.rewrite_expr(expr, input.schema())?)?
                    .build()
            }
            LogicalPlan::Aggregate {
                input,
                group_expr,
                aggr_expr,
                ..
            } => LogicalPlanBuilder::from(&self.optimize(input)?)
                .aggregate(
                    self.rewrite_expr_list(group_expr, input.schema())?,
                    self.rewrite_expr_list(aggr_expr, input.schema())?,
                )?
                .build(),
            LogicalPlan::Limit { n, input, .. } => {
                LogicalPlanBuilder::from(&self.optimize(input)?)
                    .limit(*n)?
                    .build()
            }
            LogicalPlan::Sort { input, expr, .. } => {
                LogicalPlanBuilder::from(&self.optimize(input)?)
                    .sort(self.rewrite_expr_list(expr, input.schema())?)?
                    .build()
            }
            // the following rules do not have inputs and do not need to be re-written
            LogicalPlan::TableScan { .. } => Ok(plan.clone()),
            LogicalPlan::InMemoryScan { .. } => Ok(plan.clone()),
            LogicalPlan::ParquetScan { .. } => Ok(plan.clone()),
            LogicalPlan::CsvScan { .. } => Ok(plan.clone()),
            LogicalPlan::EmptyRelation { .. } => Ok(plan.clone()),
            LogicalPlan::CreateExternalTable { .. } => Ok(plan.clone()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::execution::context::ExecutionContext;
    use crate::execution::physical_plan::csv::CsvReadOptions;
    use crate::logicalplan::{aggregate_expr, col, lit, Operator};
    use crate::test::arrow_testdata_path;
    use arrow::datatypes::{DataType, Field, Schema};

    #[test]
    fn test_all_operators() -> Result<()> {
        let testdata = arrow_testdata_path();
        let path = format!("{}/csv/aggregate_test_100.csv", testdata);

        let options = CsvReadOptions::new().schema_infer_max_records(100);
        let plan = LogicalPlanBuilder::scan_csv(&path, options, None)?
            // filter clause needs the type coercion rule applied
            .filter(col("c7").lt(&lit(5_u8)))?
            .project(vec![col("c1"), col("c2")])?
            .aggregate(
                vec![col("c1")],
                vec![aggregate_expr("SUM", col("c2"), DataType::Int64)],
            )?
            .sort(vec![col("c1")])?
            .limit(10)?
            .build()?;

        let scalar_functions = HashMap::new();
        let mut rule = TypeCoercionRule::new(&scalar_functions);
        let plan = rule.optimize(&plan)?;

        // check that the filter had a cast added
        let plan_str = format!("{:?}", plan);
        println!("{}", plan_str);
        let expected_plan_str = "Limit: 10
  Sort: #c1
    Aggregate: groupBy=[[#c1]], aggr=[[SUM(#c2)]]
      Projection: #c1, #c2
        Selection: #c7 Lt CAST(UInt8(5) AS Int64)";
        assert!(plan_str.starts_with(expected_plan_str));

        Ok(())
    }

    #[test]
    fn test_with_csv_plan() -> Result<()> {
        let testdata = arrow_testdata_path();
        let path = format!("{}/csv/aggregate_test_100.csv", testdata);

        let options = CsvReadOptions::new().schema_infer_max_records(100);
        let plan = LogicalPlanBuilder::scan_csv(&path, options, None)?
            .filter(col("c7").lt(&col("c12")))?
            .build()?;

        let scalar_functions = HashMap::new();
        let mut rule = TypeCoercionRule::new(&scalar_functions);
        let plan = rule.optimize(&plan)?;

        assert!(
            format!("{:?}", plan).starts_with("Selection: CAST(#c7 AS Float64) Lt #c12")
        );

        Ok(())
    }

    #[test]
    fn test_add_i32_i64() {
        binary_cast_test(
            DataType::Int32,
            DataType::Int64,
            "CAST(#c0 AS Int64) Plus #c1",
        );
        binary_cast_test(
            DataType::Int64,
            DataType::Int32,
            "#c0 Plus CAST(#c1 AS Int64)",
        );
    }

    #[test]
    fn test_add_f32_f64() {
        binary_cast_test(
            DataType::Float32,
            DataType::Float64,
            "CAST(#c0 AS Float64) Plus #c1",
        );
        binary_cast_test(
            DataType::Float64,
            DataType::Float32,
            "#c0 Plus CAST(#c1 AS Float64)",
        );
    }

    #[test]
    fn test_add_i32_f32() {
        binary_cast_test(
            DataType::Int32,
            DataType::Float32,
            "CAST(#c0 AS Float32) Plus #c1",
        );
        binary_cast_test(
            DataType::Float32,
            DataType::Int32,
            "#c0 Plus CAST(#c1 AS Float32)",
        );
    }

    #[test]
    fn test_add_u32_i64() {
        binary_cast_test(
            DataType::UInt32,
            DataType::Int64,
            "CAST(#c0 AS Int64) Plus #c1",
        );
        binary_cast_test(
            DataType::Int64,
            DataType::UInt32,
            "#c0 Plus CAST(#c1 AS Int64)",
        );
    }

    fn binary_cast_test(left_type: DataType, right_type: DataType, expected: &str) {
        let schema = Schema::new(vec![
            Field::new("c0", left_type, true),
            Field::new("c1", right_type, true),
        ]);

        let expr = Expr::BinaryExpr {
            left: Box::new(col("c0")),
            op: Operator::Plus,
            right: Box::new(col("c1")),
        };

        let ctx = ExecutionContext::new();
        let rule = TypeCoercionRule::new(ctx.scalar_functions());

        let expr2 = rule.rewrite_expr(&expr, &schema).unwrap();

        assert_eq!(expected, format!("{:?}", expr2));
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Mozilla
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License. You may obtain a copy of the
// License at http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

///! This module defines some core types that support find expressions: sources,
///! variables, expressions, etc.
///! These are produced as 'fuel' by the query parser, consumed by the query
///! translator and executor.
///!
///! Many of these types are defined as simple structs that are little more than
///! a richer type alias: a variable, for example, is really just a fancy kind
///! of string.
///!
///! At some point in the future, we might consider reducing copying and memory
///! usage by recasting all of these string-holding structs and enums in terms
///! of string references, with those references being slices of some parsed
///! input query string, and valid for the lifetime of that string.
///!
///! For now, for the sake of simplicity, all of these strings are heap-allocated.
///!
///! Furthermore, we might cut out some of the chaff here: each time a 'tagged'
///! type is used within an enum, we have an opportunity to simplify and use the
///! inner type directly in conjunction with matching on the enum. Before diving
///! deeply into this it's worth recognizing that this loss of 'sovereignty' is
///! a tradeoff against well-typed function signatures and other such boundaries.

extern crate edn;

use edn::{BigInt, OrderedFloat};
pub use edn::{NamespacedKeyword, PlainSymbol};

pub type SrcVarName = String;          // Do not include the required syntactic '$'.

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Variable(pub PlainSymbol);

pub trait FromValue<T> {
    fn from_value(v: &edn::Value) -> Option<T>;
}

/// If the provided EDN value is a PlainSymbol beginning with '?', return
/// it wrapped in a Variable. If not, return None.
impl FromValue<Variable> for Variable {
    fn from_value(v: &edn::Value) -> Option<Variable> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            Variable::from_symbol(s)
        } else {
            None
        }
    }
}

impl Variable {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<Variable> {
        if sym.is_var_symbol() {
            Some(Variable(sym.clone()))
        } else {
            None
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SrcVar {
    DefaultSrc,
    NamedSrc(SrcVarName),
}

impl FromValue<SrcVar> for SrcVar {
    fn from_value(v: &edn::Value) -> Option<SrcVar> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            SrcVar::from_symbol(s)
        } else {
            None
        }
    }
}

impl SrcVar {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<SrcVar> {
        if sym.is_src_symbol() {
            Some(SrcVar::NamedSrc(sym.plain_name().to_string()))
        } else {
            None
        }
    }
}

/// These are the scalar values representable in EDN.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum NonIntegerConstant {
    Boolean(bool),
    BigInteger(BigInt),
    Float(OrderedFloat<f64>),
    Text(String),
}

pub enum FnArg {
    Variable(Variable),
    SrcVar(SrcVar),
    EntidOrInteger(i64),
    Ident(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

/// e, a, tx can't be values -- no strings, no floats -- and so
/// they can only be variables, entity IDs, ident keywords, or
/// placeholders.
/// This encoding allows us to represent integers that aren't
/// entity IDs. That'll get filtered out in the context of the
/// database.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternNonValuePlace {
    Placeholder,
    Variable(Variable),
    Entid(i64),                       // Will always be +ve. See #190.
    Ident(NamespacedKeyword),
}

impl PatternNonValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_value_place(self) -> PatternValuePlace {
        match self {
            PatternNonValuePlace::Placeholder => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(x) => PatternValuePlace::Variable(x),
            PatternNonValuePlace::Entid(x)    => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(x)    => PatternValuePlace::IdentOrKeyword(x),
        }
    }

    fn to_pattern_value_place(&self) -> PatternValuePlace {
        match *self {
            PatternNonValuePlace::Placeholder     => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(ref x) => PatternValuePlace::Variable(x.clone()),
            PatternNonValuePlace::Entid(x)        => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(ref x)    => PatternValuePlace::IdentOrKeyword(x.clone()),
        }
    }
}

impl FromValue<PatternNonValuePlace> for PatternNonValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternNonValuePlace> {
        match v {
            &edn::Value::Integer(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            &edn::Value::PlainSymbol(ref x) => if x.0.as_str() == "_" {
                Some(PatternNonValuePlace::Placeholder)
            } else {
                if let Some(v) = Variable::from_symbol(x) {
                    Some(PatternNonValuePlace::Variable(v))
                } else {
                    None
                }
            },
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternNonValuePlace::Ident(x.clone())),
            _ => None,
        }
    }
}

/// The `v` part of a pattern can be much broader: it can represent
/// integers that aren't entity IDs (particularly negative integers),
/// strings, and all the rest. We group those under `Constant`.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternValuePlace {
    Placeholder,
    Variable(Variable),
    EntidOrInteger(i64),
    IdentOrKeyword(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

impl FromValue<PatternValuePlace> for PatternValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternValuePlace> {
        match v {
            &edn::Value::Integer(x) =>
                Some(PatternValuePlace::EntidOrInteger(x)),
            &edn::Value::PlainSymbol(ref x) if x.0.as_str() == "_" =>
                Some(PatternValuePlace::Placeholder),
            &edn::Value::PlainSymbol(ref x) if x.is_var_symbol() =>
                Some(PatternValuePlace::Variable(Variable(x.clone()))),
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternValuePlace::IdentOrKeyword(x.clone())),
            &edn::Value::Boolean(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Boolean(x))),
            &edn::Value::Float(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Float(x))),
            &edn::Value::BigInteger(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::BigInteger(x.clone()))),
            &edn::Value::Text(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Text(x.clone()))),
            _ => None,
        }
    }
}

impl PatternValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_non_value_place(self) -> Option<PatternNonValuePlace> {
        match self {
            PatternValuePlace::Placeholder       => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(x)       => Some(PatternNonValuePlace::Variable(x)),
            PatternValuePlace::EntidOrInteger(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(x) => Some(PatternNonValuePlace::Ident(x)),
            PatternValuePlace::Constant(_)       => None,
        }
    }

    fn to_pattern_non_value_place(&self) -> Option<PatternNonValuePlace> {
        match *self {
            PatternValuePlace::Placeholder           => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(ref x)       => Some(PatternNonValuePlace::Variable(x.clone())),
            PatternValuePlace::EntidOrInteger(x)     => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(ref x) => Some(PatternNonValuePlace::Ident(x.clone())),
            PatternValuePlace::Constant(_)           => None,
        }
    }
}

/*
pub enum PullPattern {
    Constant(Constant),
    Variable(Variable),
}

pub struct Pull {
    pub src: SrcVar,
    pub var: Variable,
    pub pattern: PullPattern,      // Constant, variable, or plain variable.
}
*/

/*
pub struct Aggregate {
    pub fn_name: String,
    pub args: Vec<FnArg>,
}
*/

#[derive(Clone,Debug,Eq,PartialEq)]
pub enum Element {
    Variable(Variable),
    // Aggregate(Aggregate),   // TODO
    // Pull(Pull),             // TODO
}

/// A definition of the first part of a find query: the
/// `[:find ?foo ?bar…]` bit.
///
/// There are four different kinds of find specs, allowing you to query for
/// a single value, a collection of values from different entities, a single
/// tuple (relation), or a collection of tuples.
///
/// Examples:
///
/// ```rust
/// # extern crate edn;
/// # extern crate mentat_query;
/// # use edn::PlainSymbol;
/// # use mentat_query::{Element, FindSpec, Variable};
///
/// # fn main() {
///
///   let elements = vec![
///     Element::Variable(Variable(PlainSymbol("?foo".to_string()))),
///     Element::Variable(Variable(PlainSymbol("?bar".to_string()))),
///   ];
///   let rel = FindSpec::FindRel(elements);
///
///   if let FindSpec::FindRel(elements) = rel {
///     assert_eq!(2, elements.len());
///   }
///
/// # }
/// ```
///
#[derive(Clone,Debug,Eq,PartialEq)]
pub enum FindSpec {
    /// Returns an array of arrays.
    FindRel(Vec<Element>),

    /// Returns an array of scalars, usually homogeneous.
    /// This is equivalent to mapping over the results of a `FindRel`,
    /// returning the first value of each.
    FindColl(Element),

    /// Returns a single tuple: a heterogeneous array of scalars. Equivalent to
    /// taking the first result from a `FindRel`.
    FindTuple(Vec<Element>),

    /// Returns a single scalar value. Equivalent to taking the first result
    /// from a `FindColl`.
    FindScalar(Element),
}

/// Returns true if the provided `FindSpec` returns at most one result.
pub fn is_unit_limited(spec: &FindSpec) -> bool {
    match spec {
        &FindSpec::FindScalar(..) => true,
        &FindSpec::FindTuple(..)  => true,
        &FindSpec::FindRel(..)    => false,
        &FindSpec::FindColl(..)   => false,
    }
}

/// Returns true if the provided `FindSpec` cares about distinct results.
///
/// I use the words "cares about" because find is generally defined in terms of producing distinct
/// results at the Datalog level.
///
/// Two of the find specs (scalar and tuple) produce only a single result. Those don't need to be
/// run with `SELECT DISTINCT`, because we're only consuming a single result. Those queries will be
/// run with `LIMIT 1`.
///
/// Additionally, some projections cannot produce duplicate results: `[:find (max ?x) …]`, for
/// example.
///
/// This function gives us the hook to add that logic when we're ready.
///
/// Beyond this, `DISTINCT` is not always needed. For example, in some kinds of accumulation or
/// sampling projections we might not need to do it at the SQL level because we're consuming into
/// a dupe-eliminating data structure like a Set, or we know that a particular query cannot produce
/// duplicate results.
pub fn requires_distinct(spec: &FindSpec) -> bool {
    return !is_unit_limited(spec);
}

// Note that the "implicit blank" rule applies.
// A pattern with a reversed attribute — :foo/_bar — is reversed
// at the point of parsing. These `Pattern` instances only represent
// one direction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Pattern {
    pub source: Option<SrcVar>,
    pub entity: PatternNonValuePlace,
    pub attribute: PatternNonValuePlace,
    pub value: PatternValuePlace,
    pub tx: PatternNonValuePlace,
}

impl Pattern {
    pub fn new(src: Option<SrcVar>,
               e: PatternNonValuePlace,
               a: PatternNonValuePlace,
               v: PatternValuePlace,
               tx: PatternNonValuePlace) -> Option<Pattern> {
        let aa = a.clone();       // Too tired of fighting borrow scope for now.
        if let PatternNonValuePlace::Ident(ref k) = aa {
            if k.is_backward() {
                // e and v have different types; we must convert them.
                // Not every parseable value is suitable for the entity field!
                // As such, this is a failable constructor.
                let e_v = e.to_pattern_value_place();
                if let Some(v_e) = v.to_pattern_non_value_place() {
                    return Some(Pattern {
                        source: src,
                        entity: v_e,
                        attribute: PatternNonValuePlace::Ident(k.to_reversed()),
                        value: e_v,
                        tx: tx,
                    });
                } else {
                    return None;
                }
            }
        }
        Some(Pattern {
            source: src,
            entity: e,
            attribute: a,
            value: v,
            tx: tx,
        })
    }
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum WhereClause {
    Not,
    NotJoin,
    Or,
    OrJoin,
    Pred,
    WhereFn,
    RuleExpr,
    Pattern(Pattern),
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FindQuery {
    pub find_spec: FindSpec,
    pub default_source: SrcVar,
    pub with: Vec<Variable>,
    pub in_vars: Vec<Variable>,
    pub in_sources: Vec<SrcVar>,
    pub where_clauses: Vec<WhereClause>,
    // TODO: in_rules;
}
<commit_msg>Derive PartialOrd and Ord for Variable. r=jsantell<commit_after>// Copyright 2016 Mozilla
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License. You may obtain a copy of the
// License at http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

///! This module defines some core types that support find expressions: sources,
///! variables, expressions, etc.
///! These are produced as 'fuel' by the query parser, consumed by the query
///! translator and executor.
///!
///! Many of these types are defined as simple structs that are little more than
///! a richer type alias: a variable, for example, is really just a fancy kind
///! of string.
///!
///! At some point in the future, we might consider reducing copying and memory
///! usage by recasting all of these string-holding structs and enums in terms
///! of string references, with those references being slices of some parsed
///! input query string, and valid for the lifetime of that string.
///!
///! For now, for the sake of simplicity, all of these strings are heap-allocated.
///!
///! Furthermore, we might cut out some of the chaff here: each time a 'tagged'
///! type is used within an enum, we have an opportunity to simplify and use the
///! inner type directly in conjunction with matching on the enum. Before diving
///! deeply into this it's worth recognizing that this loss of 'sovereignty' is
///! a tradeoff against well-typed function signatures and other such boundaries.

extern crate edn;

use edn::{BigInt, OrderedFloat};
pub use edn::{NamespacedKeyword, PlainSymbol};

pub type SrcVarName = String;          // Do not include the required syntactic '$'.

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Variable(pub PlainSymbol);

pub trait FromValue<T> {
    fn from_value(v: &edn::Value) -> Option<T>;
}

/// If the provided EDN value is a PlainSymbol beginning with '?', return
/// it wrapped in a Variable. If not, return None.
impl FromValue<Variable> for Variable {
    fn from_value(v: &edn::Value) -> Option<Variable> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            Variable::from_symbol(s)
        } else {
            None
        }
    }
}

impl Variable {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<Variable> {
        if sym.is_var_symbol() {
            Some(Variable(sym.clone()))
        } else {
            None
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SrcVar {
    DefaultSrc,
    NamedSrc(SrcVarName),
}

impl FromValue<SrcVar> for SrcVar {
    fn from_value(v: &edn::Value) -> Option<SrcVar> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            SrcVar::from_symbol(s)
        } else {
            None
        }
    }
}

impl SrcVar {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<SrcVar> {
        if sym.is_src_symbol() {
            Some(SrcVar::NamedSrc(sym.plain_name().to_string()))
        } else {
            None
        }
    }
}

/// These are the scalar values representable in EDN.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum NonIntegerConstant {
    Boolean(bool),
    BigInteger(BigInt),
    Float(OrderedFloat<f64>),
    Text(String),
}

pub enum FnArg {
    Variable(Variable),
    SrcVar(SrcVar),
    EntidOrInteger(i64),
    Ident(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

/// e, a, tx can't be values -- no strings, no floats -- and so
/// they can only be variables, entity IDs, ident keywords, or
/// placeholders.
/// This encoding allows us to represent integers that aren't
/// entity IDs. That'll get filtered out in the context of the
/// database.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternNonValuePlace {
    Placeholder,
    Variable(Variable),
    Entid(i64),                       // Will always be +ve. See #190.
    Ident(NamespacedKeyword),
}

impl PatternNonValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_value_place(self) -> PatternValuePlace {
        match self {
            PatternNonValuePlace::Placeholder => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(x) => PatternValuePlace::Variable(x),
            PatternNonValuePlace::Entid(x)    => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(x)    => PatternValuePlace::IdentOrKeyword(x),
        }
    }

    fn to_pattern_value_place(&self) -> PatternValuePlace {
        match *self {
            PatternNonValuePlace::Placeholder     => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(ref x) => PatternValuePlace::Variable(x.clone()),
            PatternNonValuePlace::Entid(x)        => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(ref x)    => PatternValuePlace::IdentOrKeyword(x.clone()),
        }
    }
}

impl FromValue<PatternNonValuePlace> for PatternNonValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternNonValuePlace> {
        match v {
            &edn::Value::Integer(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            &edn::Value::PlainSymbol(ref x) => if x.0.as_str() == "_" {
                Some(PatternNonValuePlace::Placeholder)
            } else {
                if let Some(v) = Variable::from_symbol(x) {
                    Some(PatternNonValuePlace::Variable(v))
                } else {
                    None
                }
            },
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternNonValuePlace::Ident(x.clone())),
            _ => None,
        }
    }
}

/// The `v` part of a pattern can be much broader: it can represent
/// integers that aren't entity IDs (particularly negative integers),
/// strings, and all the rest. We group those under `Constant`.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternValuePlace {
    Placeholder,
    Variable(Variable),
    EntidOrInteger(i64),
    IdentOrKeyword(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

impl FromValue<PatternValuePlace> for PatternValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternValuePlace> {
        match v {
            &edn::Value::Integer(x) =>
                Some(PatternValuePlace::EntidOrInteger(x)),
            &edn::Value::PlainSymbol(ref x) if x.0.as_str() == "_" =>
                Some(PatternValuePlace::Placeholder),
            &edn::Value::PlainSymbol(ref x) if x.is_var_symbol() =>
                Some(PatternValuePlace::Variable(Variable(x.clone()))),
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternValuePlace::IdentOrKeyword(x.clone())),
            &edn::Value::Boolean(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Boolean(x))),
            &edn::Value::Float(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Float(x))),
            &edn::Value::BigInteger(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::BigInteger(x.clone()))),
            &edn::Value::Text(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Text(x.clone()))),
            _ => None,
        }
    }
}

impl PatternValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_non_value_place(self) -> Option<PatternNonValuePlace> {
        match self {
            PatternValuePlace::Placeholder       => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(x)       => Some(PatternNonValuePlace::Variable(x)),
            PatternValuePlace::EntidOrInteger(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(x) => Some(PatternNonValuePlace::Ident(x)),
            PatternValuePlace::Constant(_)       => None,
        }
    }

    fn to_pattern_non_value_place(&self) -> Option<PatternNonValuePlace> {
        match *self {
            PatternValuePlace::Placeholder           => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(ref x)       => Some(PatternNonValuePlace::Variable(x.clone())),
            PatternValuePlace::EntidOrInteger(x)     => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(ref x) => Some(PatternNonValuePlace::Ident(x.clone())),
            PatternValuePlace::Constant(_)           => None,
        }
    }
}

/*
pub enum PullPattern {
    Constant(Constant),
    Variable(Variable),
}

pub struct Pull {
    pub src: SrcVar,
    pub var: Variable,
    pub pattern: PullPattern,      // Constant, variable, or plain variable.
}
*/

/*
pub struct Aggregate {
    pub fn_name: String,
    pub args: Vec<FnArg>,
}
*/

#[derive(Clone,Debug,Eq,PartialEq)]
pub enum Element {
    Variable(Variable),
    // Aggregate(Aggregate),   // TODO
    // Pull(Pull),             // TODO
}

/// A definition of the first part of a find query: the
/// `[:find ?foo ?bar…]` bit.
///
/// There are four different kinds of find specs, allowing you to query for
/// a single value, a collection of values from different entities, a single
/// tuple (relation), or a collection of tuples.
///
/// Examples:
///
/// ```rust
/// # extern crate edn;
/// # extern crate mentat_query;
/// # use edn::PlainSymbol;
/// # use mentat_query::{Element, FindSpec, Variable};
///
/// # fn main() {
///
///   let elements = vec![
///     Element::Variable(Variable(PlainSymbol("?foo".to_string()))),
///     Element::Variable(Variable(PlainSymbol("?bar".to_string()))),
///   ];
///   let rel = FindSpec::FindRel(elements);
///
///   if let FindSpec::FindRel(elements) = rel {
///     assert_eq!(2, elements.len());
///   }
///
/// # }
/// ```
///
#[derive(Clone,Debug,Eq,PartialEq)]
pub enum FindSpec {
    /// Returns an array of arrays.
    FindRel(Vec<Element>),

    /// Returns an array of scalars, usually homogeneous.
    /// This is equivalent to mapping over the results of a `FindRel`,
    /// returning the first value of each.
    FindColl(Element),

    /// Returns a single tuple: a heterogeneous array of scalars. Equivalent to
    /// taking the first result from a `FindRel`.
    FindTuple(Vec<Element>),

    /// Returns a single scalar value. Equivalent to taking the first result
    /// from a `FindColl`.
    FindScalar(Element),
}

/// Returns true if the provided `FindSpec` returns at most one result.
pub fn is_unit_limited(spec: &FindSpec) -> bool {
    match spec {
        &FindSpec::FindScalar(..) => true,
        &FindSpec::FindTuple(..)  => true,
        &FindSpec::FindRel(..)    => false,
        &FindSpec::FindColl(..)   => false,
    }
}

/// Returns true if the provided `FindSpec` cares about distinct results.
///
/// I use the words "cares about" because find is generally defined in terms of producing distinct
/// results at the Datalog level.
///
/// Two of the find specs (scalar and tuple) produce only a single result. Those don't need to be
/// run with `SELECT DISTINCT`, because we're only consuming a single result. Those queries will be
/// run with `LIMIT 1`.
///
/// Additionally, some projections cannot produce duplicate results: `[:find (max ?x) …]`, for
/// example.
///
/// This function gives us the hook to add that logic when we're ready.
///
/// Beyond this, `DISTINCT` is not always needed. For example, in some kinds of accumulation or
/// sampling projections we might not need to do it at the SQL level because we're consuming into
/// a dupe-eliminating data structure like a Set, or we know that a particular query cannot produce
/// duplicate results.
pub fn requires_distinct(spec: &FindSpec) -> bool {
    return !is_unit_limited(spec);
}

// Note that the "implicit blank" rule applies.
// A pattern with a reversed attribute — :foo/_bar — is reversed
// at the point of parsing. These `Pattern` instances only represent
// one direction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Pattern {
    pub source: Option<SrcVar>,
    pub entity: PatternNonValuePlace,
    pub attribute: PatternNonValuePlace,
    pub value: PatternValuePlace,
    pub tx: PatternNonValuePlace,
}

impl Pattern {
    pub fn new(src: Option<SrcVar>,
               e: PatternNonValuePlace,
               a: PatternNonValuePlace,
               v: PatternValuePlace,
               tx: PatternNonValuePlace) -> Option<Pattern> {
        let aa = a.clone();       // Too tired of fighting borrow scope for now.
        if let PatternNonValuePlace::Ident(ref k) = aa {
            if k.is_backward() {
                // e and v have different types; we must convert them.
                // Not every parseable value is suitable for the entity field!
                // As such, this is a failable constructor.
                let e_v = e.to_pattern_value_place();
                if let Some(v_e) = v.to_pattern_non_value_place() {
                    return Some(Pattern {
                        source: src,
                        entity: v_e,
                        attribute: PatternNonValuePlace::Ident(k.to_reversed()),
                        value: e_v,
                        tx: tx,
                    });
                } else {
                    return None;
                }
            }
        }
        Some(Pattern {
            source: src,
            entity: e,
            attribute: a,
            value: v,
            tx: tx,
        })
    }
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum WhereClause {
    Not,
    NotJoin,
    Or,
    OrJoin,
    Pred,
    WhereFn,
    RuleExpr,
    Pattern(Pattern),
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FindQuery {
    pub find_spec: FindSpec,
    pub default_source: SrcVar,
    pub with: Vec<Variable>,
    pub in_vars: Vec<Variable>,
    pub in_sources: Vec<SrcVar>,
    pub where_clauses: Vec<WhereClause>,
    // TODO: in_rules;
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

#![feature(plugin, use_extern_macros)]
#![plugin(tarpc_plugins)]

extern crate env_logger;
extern crate futures;
#[macro_use]
extern crate tarpc;
extern crate tokio_core;

use futures::{Future, future};
use publisher::FutureServiceExt as PublisherExt;
use std::cell::RefCell;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::rc::Rc;
use std::thread;
use std::time::Duration;
use subscriber::FutureServiceExt as SubscriberExt;
use tarpc::future::{client, server};
use tarpc::future::client::ClientExt;
use tarpc::util::{FirstSocketAddr, Message, Never};
use tokio_core::reactor;

pub mod subscriber {
    service! {
        rpc receive(message: String);
    }
}

pub mod publisher {
    use std::net::SocketAddr;
    use tarpc::util::Message;

    service! {
        rpc broadcast(message: String);
        rpc subscribe(id: u32, address: SocketAddr) | Message;
        rpc unsubscribe(id: u32);
    }
}

#[derive(Clone, Debug)]
struct Subscriber {
    id: u32,
}

impl subscriber::FutureService for Subscriber {
    type ReceiveFut = Result<(), Never>;

    fn receive(&self, message: String) -> Self::ReceiveFut {
        println!("{} received message: {}", self.id, message);
        Ok(())
    }
}

impl Subscriber {
    fn listen(id: u32, handle: &reactor::Handle, options: server::Options) -> server::Handle {
        let (server_handle, server) = Subscriber { id: id }
            .listen("localhost:0".first_socket_addr(), handle, options)
            .unwrap();
        handle.spawn(server);
        server_handle
    }
}

#[derive(Clone, Debug)]
struct Publisher {
    clients: Rc<RefCell<HashMap<u32, subscriber::FutureClient>>>,
}

impl Publisher {
    fn new() -> Publisher {
        Publisher {
            clients: Rc::new(RefCell::new(HashMap::new())),
        }
    }
}

impl publisher::FutureService for Publisher {
    type BroadcastFut = Box<Future<Item = (), Error = Never>>;

    fn broadcast(&self, message: String) -> Self::BroadcastFut {
        let acks = self.clients
                       .borrow()
                       .values()
                       .map(move |client| client.receive(message.clone())
                           // Ignore failing subscribers. In a real pubsub,
                           // you'd want to continually retry until subscribers
                           // ack.
                           .then(|_| Ok(())))
                       // Collect to a vec to end the borrow on `self.clients`.
                       .collect::<Vec<_>>();
        Box::new(future::join_all(acks).map(|_| ()))
    }

    type SubscribeFut = Box<Future<Item = (), Error = Message>>;

    fn subscribe(&self, id: u32, address: SocketAddr) -> Self::SubscribeFut {
        let clients = self.clients.clone();
        Box::new(
            subscriber::FutureClient::connect(address, client::Options::default())
                .map(move |subscriber| {
                    println!("Subscribing {}.", id);
                    clients.borrow_mut().insert(id, subscriber);
                    ()
                })
                .map_err(|e| e.to_string().into()),
        )
    }

    type UnsubscribeFut = Box<Future<Item = (), Error = Never>>;

    fn unsubscribe(&self, id: u32) -> Self::UnsubscribeFut {
        println!("Unsubscribing {}", id);
        self.clients.borrow_mut().remove(&id).unwrap();
        futures::finished(()).boxed()
    }
}

fn main() {
    let _ = env_logger::init();
    let mut reactor = reactor::Core::new().unwrap();
    let (publisher_handle, server) = Publisher::new()
        .listen(
            "localhost:0".first_socket_addr(),
            &reactor.handle(),
            server::Options::default(),
        )
        .unwrap();
    reactor.handle().spawn(server);

    let subscriber1 = Subscriber::listen(0, &reactor.handle(), server::Options::default());
    let subscriber2 = Subscriber::listen(1, &reactor.handle(), server::Options::default());

    let publisher = reactor
        .run(publisher::FutureClient::connect(
            publisher_handle.addr(),
            client::Options::default(),
        ))
        .unwrap();
    reactor
        .run(
            publisher
                .subscribe(0, subscriber1.addr())
                .and_then(|_| publisher.subscribe(1, subscriber2.addr()))
                .map_err(|e| panic!(e))
                .and_then(|_| {
                    println!("Broadcasting...");
                    publisher.broadcast("hello to all".to_string())
                })
                .and_then(|_| publisher.unsubscribe(1))
                .and_then(|_| publisher.broadcast("hi again".to_string())),
        )
        .unwrap();
    thread::sleep(Duration::from_millis(300));
}
<commit_msg>clippy: using '.clone()' on a ref-counted pointer<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

#![feature(plugin, use_extern_macros)]
#![plugin(tarpc_plugins)]

extern crate env_logger;
extern crate futures;
#[macro_use]
extern crate tarpc;
extern crate tokio_core;

use futures::{Future, future};
use publisher::FutureServiceExt as PublisherExt;
use std::cell::RefCell;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::rc::Rc;
use std::thread;
use std::time::Duration;
use subscriber::FutureServiceExt as SubscriberExt;
use tarpc::future::{client, server};
use tarpc::future::client::ClientExt;
use tarpc::util::{FirstSocketAddr, Message, Never};
use tokio_core::reactor;

pub mod subscriber {
    service! {
        rpc receive(message: String);
    }
}

pub mod publisher {
    use std::net::SocketAddr;
    use tarpc::util::Message;

    service! {
        rpc broadcast(message: String);
        rpc subscribe(id: u32, address: SocketAddr) | Message;
        rpc unsubscribe(id: u32);
    }
}

#[derive(Clone, Debug)]
struct Subscriber {
    id: u32,
}

impl subscriber::FutureService for Subscriber {
    type ReceiveFut = Result<(), Never>;

    fn receive(&self, message: String) -> Self::ReceiveFut {
        println!("{} received message: {}", self.id, message);
        Ok(())
    }
}

impl Subscriber {
    fn listen(id: u32, handle: &reactor::Handle, options: server::Options) -> server::Handle {
        let (server_handle, server) = Subscriber { id: id }
            .listen("localhost:0".first_socket_addr(), handle, options)
            .unwrap();
        handle.spawn(server);
        server_handle
    }
}

#[derive(Clone, Debug)]
struct Publisher {
    clients: Rc<RefCell<HashMap<u32, subscriber::FutureClient>>>,
}

impl Publisher {
    fn new() -> Publisher {
        Publisher {
            clients: Rc::new(RefCell::new(HashMap::new())),
        }
    }
}

impl publisher::FutureService for Publisher {
    type BroadcastFut = Box<Future<Item = (), Error = Never>>;

    fn broadcast(&self, message: String) -> Self::BroadcastFut {
        let acks = self.clients
                       .borrow()
                       .values()
                       .map(move |client| client.receive(message.clone())
                           // Ignore failing subscribers. In a real pubsub,
                           // you'd want to continually retry until subscribers
                           // ack.
                           .then(|_| Ok(())))
                       // Collect to a vec to end the borrow on `self.clients`.
                       .collect::<Vec<_>>();
        Box::new(future::join_all(acks).map(|_| ()))
    }

    type SubscribeFut = Box<Future<Item = (), Error = Message>>;

    fn subscribe(&self, id: u32, address: SocketAddr) -> Self::SubscribeFut {
        let clients = Rc::clone(&self.clients);
        Box::new(
            subscriber::FutureClient::connect(address, client::Options::default())
                .map(move |subscriber| {
                    println!("Subscribing {}.", id);
                    clients.borrow_mut().insert(id, subscriber);
                    ()
                })
                .map_err(|e| e.to_string().into()),
        )
    }

    type UnsubscribeFut = Box<Future<Item = (), Error = Never>>;

    fn unsubscribe(&self, id: u32) -> Self::UnsubscribeFut {
        println!("Unsubscribing {}", id);
        self.clients.borrow_mut().remove(&id).unwrap();
        futures::finished(()).boxed()
    }
}

fn main() {
    let _ = env_logger::init();
    let mut reactor = reactor::Core::new().unwrap();
    let (publisher_handle, server) = Publisher::new()
        .listen(
            "localhost:0".first_socket_addr(),
            &reactor.handle(),
            server::Options::default(),
        )
        .unwrap();
    reactor.handle().spawn(server);

    let subscriber1 = Subscriber::listen(0, &reactor.handle(), server::Options::default());
    let subscriber2 = Subscriber::listen(1, &reactor.handle(), server::Options::default());

    let publisher = reactor
        .run(publisher::FutureClient::connect(
            publisher_handle.addr(),
            client::Options::default(),
        ))
        .unwrap();
    reactor
        .run(
            publisher
                .subscribe(0, subscriber1.addr())
                .and_then(|_| publisher.subscribe(1, subscriber2.addr()))
                .map_err(|e| panic!(e))
                .and_then(|_| {
                    println!("Broadcasting...");
                    publisher.broadcast("hello to all".to_string())
                })
                .and_then(|_| publisher.unsubscribe(1))
                .and_then(|_| publisher.broadcast("hi again".to_string())),
        )
        .unwrap();
    thread::sleep(Duration::from_millis(300));
}
<|endoftext|>