<commit_before>use std::ffi::CString;
use std::os::unix::ffi::OsStrExt;

use c::FileType;
use c::XmpPacketInfo;

use crate::error::Error;
use crate::xmp::Xmp;
use crate::xmpstring::XmpString;
use crate::Result;

bitflags! {
    /// Flag options for opening files.
    #[derive(Default)]
    pub struct OpenFlags: u32 {
        /// No open option
        const NONE = 0x0000_0000;
        /// Open for read-only access.
        const READ = 0x0000_0001;
        /// Open for reading and writing.
        const FOR_UPDATE = 0x0000_0002;
        /// Only the XMP is wanted, allows space/time optimizations.
        const ONLY_XMP = 0x0000_0004;
        /// Cache thumbnail if possible, GetThumbnail will be called.
        const CACHE_TNAIL = 0x0000_0008;
        /// Be strict about locating XMP and reconciling with other forms.
        const STRICTLY = 0x0000_0010;
        /// Require the use of a smart handler.
        const USE_SMART_HANDLER = 0x0000_0020;
        /// Force packet scanning, don't use a smart handler.
        const USE_PACKET_SCANNING = 0x0000_0040;
        /// Only packet scan files "known" to need scanning.
        const LIMITED_SCANNING = 0x0000_0080;
        /// Attempt to repair a file opened for update, default is to not open (throw an exception).
        const REPAIR_FILE = 0x0000_0100;
        /// Optimize MPEG4 to support stream when updating This can take some time
        const OPTIMIZE_FILE_LAYOUT = 0x0000_0200;
        /// Set if calling from background thread.
        const IN_BACKGROUND = 0x1000_0000;
    }
}

bitflags! {
    /// Flag options to close files.
    #[derive(Default)]
    pub struct CloseFlags: u32 {
        /// No close option
        const NONE = 0x0000;
        /// Write into a temporary file and swap for crash safety.
        const SAFE_UPDATE = 0x0001;
    }
}

bitflags! {
    /// Result flag for file / format infos.
    #[derive(Default)]
    pub struct FormatOptionFlags: u32 {
        const NONE = 0;
        /// Can inject first-time XMP into an existing file.
        const CAN_INJECT_XMP = 0x0000_0001;
        /// Can expand XMP or other metadata in an existing file.
        const CAN_EXPAND = 0x0000_0002;
        /// Can copy one file to another, writing new metadata.
        const CAN_REWRITE = 0x0000_0004;
        /// Can expand, but prefers in-place update.
        const PREFERS_IN_PLACE = 0x0000_0008;
        /// Supports reconciliation between XMP and other forms.
        const CAN_RECONCILE = 0x0000_0010;
        /// Allows access to just the XMP, ignoring other forms.
        const ALLOWS_ONLY_XMP = 0x0000_0020;
        /// File handler returns raw XMP packet information.
        const RETURNS_RAW_PACKET = 0x0000_0040;
        /// The file handler does the file open and close.
        const HANDLER_OWNS_FILE = 0x0000_0100;
        /// The file handler allows crash-safe file updates.
        const ALLOW_SAFE_UPDATE = 0x0000_0200;
        /// The file format needs the XMP packet to be read-only.
        const NEEDS_READONLY_PACKET = 0x0000_0400;
        /// The file handler uses a "sidecar" file for the XMP.
        const USE_SIDECAR_XMP = 0x0000_0800;
        /// The format is folder oriented, for example the P2 video format.
        const FOLDER_BASED_FORMAT = 0x0000_1000;
    }
}

/// Wrap an XmpPacketInfo
pub struct PacketInfo(pub XmpPacketInfo);

impl Default for PacketInfo {
    fn default() -> PacketInfo {
        PacketInfo(XmpPacketInfo {
            offset: 0,
            length: 0,
            pad_size: 0,
            char_form: 0,
            writeable: false,
            has_wrapper: false,
            pad: 0,
        })
    }
}

pub struct XmpFile(*mut c::XmpFile);

impl Default for XmpFile {
    fn default() -> XmpFile {
        super::init();

        XmpFile(unsafe { c::xmp_files_new() })
    }
}

impl XmpFile {
    /// Create new XmpFile
    pub fn new() -> XmpFile {
        XmpFile::default()
    }

    /// Create and open a new XmpFile
    /// Equivalent to calling new then open.
    /// Return Err in case of failure
    pub fn new_from_file<P: AsRef<std::ffi::OsStr>>(
        path: P,
        options: OpenFlags,
    ) -> Result<XmpFile> {
        super::init();

        let path = path.as_ref().as_bytes();
        let pp = CString::new(path).unwrap();
        let ptr = unsafe { c::xmp_files_open_new(pp.as_ptr(), options.bits()) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(XmpFile(ptr))
    }

    /// Open an XmpFile. Usually called after new.
    pub fn open<P: AsRef<std::ffi::OsStr>>(&mut self, path: P, options: OpenFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        let path = path.as_ref().as_bytes();
        let pp = CString::new(path).unwrap();
        if unsafe { c::xmp_files_open(self.0, pp.as_ptr(), options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Close the XmpFile
    pub fn close(&mut self, options: CloseFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_close(self.0, options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if native pointer is null
    pub fn is_null(&self) -> bool {
        self.0.is_null()
    }

    /// Get a new Xmp fronm the currently open file
    pub fn get_new_xmp(&self) -> Result<Xmp> {
        let ptr = unsafe { c::xmp_files_get_new_xmp(self.0) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(Xmp::from(ptr))
    }

    /// Get the xmp data an Xmp.
    pub fn get_xmp(&self) -> Result<Xmp> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        let mut xmp = Xmp::default();
        if unsafe { c::xmp_files_get_xmp(self.0, xmp.as_mut_ptr()) } {
            Ok(xmp)
        } else {
            Err(crate::get_error())
        }
    }

    /// Get the xmp packet as a string.
    pub fn get_xmp_xmpstring(&self) -> Result<(XmpString, PacketInfo)> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        let mut packet = XmpString::new();
        let mut info = PacketInfo::default();
        if unsafe {
            c::xmp_files_get_xmp_xmpstring(
                self.0,
                packet.as_mut_ptr(),
                &mut info.0 as *mut XmpPacketInfo,
            )
        } {
            Ok((packet, info))
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if it can put the Xmp into the XmpFile.
    pub fn can_put_xmp(&self, xmp: &Xmp) -> bool {
        if self.is_null() || xmp.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp(self.0, xmp.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_xmpstring(&self, xmp_packet: &XmpString) -> bool {
        if self.is_null() || xmp_packet.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp_xmpstring(self.0, xmp_packet.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_str(&self, xmp_packet: &str) -> bool {
        if self.is_null() {
            return false;
        }
        let pp = CString::new(xmp_packet).unwrap();
        unsafe { c::xmp_files_can_put_xmp_cstr(self.0, pp.as_ptr(), xmp_packet.len()) }
    }

    /// Put the Xmp into the XmpFile
    pub fn put_xmp(&mut self, xmp: &Xmp) -> Result<()> {
        if self.is_null() || xmp.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_put_xmp(self.0, xmp.as_ptr()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Get info from the XmpFile.
    pub fn get_file_info(
        &self,
        file_path: &mut String,
        options: &mut OpenFlags,
        format: &mut FileType,
        handler_flags: &mut FormatOptionFlags,
    ) -> bool {
        if self.is_null() {
            return false;
        }
        let mut s = XmpString::new();

        let mut raw_options: u32 = 0;
        let mut raw_handler_flags: u32 = 0;
        let result = unsafe {
            c::xmp_files_get_file_info(
                self.0,
                s.as_mut_ptr(),
                &mut raw_options,
                format,
                &mut raw_handler_flags,
            )
        };
        *options = OpenFlags::from_bits(raw_options).unwrap_or_default();
        *handler_flags = FormatOptionFlags::from_bits(raw_handler_flags).unwrap_or_default();

        file_path.push_str(&s.to_string());
        result
    }

    /// Check the file format for the specified path
    pub fn check_file_format<P: AsRef<std::ffi::OsStr>>(path: P) -> FileType {
        let path = path.as_ref().as_bytes();
        let pp = CString::new(path).unwrap();
        unsafe { c::xmp_files_check_file_format(pp.as_ptr()) }
    }

    /// Get FormatOptions for the FileType
    pub fn get_format_info(format: FileType) -> Result<FormatOptionFlags> {
        let mut raw_options: u32 = 0;
        if unsafe { c::xmp_files_get_format_info(format, &mut raw_options) } {
            Ok(FormatOptionFlags::from_bits(raw_options).unwrap_or_default())
        } else {
            Err(crate::get_error())
        }
    }
}

impl Drop for XmpFile {
    /// Drop the XmpFile.
    fn drop(&mut self) {
        if !self.is_null() {
            unsafe { c::xmp_files_free(self.0) };
        }
    }
}

#[test]
fn it_works() {
    use crate::error::Error;

    let inited = crate::init();

    assert!(inited);

    let xf = XmpFile::new();
    assert!(!xf.is_null());

    assert!(crate::get_error() == Error::from(c::XmpError::Unknown));
}
<commit_msg>Restore API for get_xmp()<commit_after>use std::ffi::CString;
use std::os::unix::ffi::OsStrExt;

use c::FileType;
use c::XmpPacketInfo;

use crate::error::Error;
use crate::xmp::Xmp;
use crate::xmpstring::XmpString;
use crate::Result;

bitflags! {
    /// Flag options for opening files.
    #[derive(Default)]
    pub struct OpenFlags: u32 {
        /// No open option
        const NONE = 0x0000_0000;
        /// Open for read-only access.
        const READ = 0x0000_0001;
        /// Open for reading and writing.
        const FOR_UPDATE = 0x0000_0002;
        /// Only the XMP is wanted, allows space/time optimizations.
        const ONLY_XMP = 0x0000_0004;
        /// Cache thumbnail if possible, GetThumbnail will be called.
        const CACHE_TNAIL = 0x0000_0008;
        /// Be strict about locating XMP and reconciling with other forms.
        const STRICTLY = 0x0000_0010;
        /// Require the use of a smart handler.
        const USE_SMART_HANDLER = 0x0000_0020;
        /// Force packet scanning, don't use a smart handler.
        const USE_PACKET_SCANNING = 0x0000_0040;
        /// Only packet scan files "known" to need scanning.
        const LIMITED_SCANNING = 0x0000_0080;
        /// Attempt to repair a file opened for update, default is to not open (throw an exception).
        const REPAIR_FILE = 0x0000_0100;
        /// Optimize MPEG4 to support stream when updating This can take some time
        const OPTIMIZE_FILE_LAYOUT = 0x0000_0200;
        /// Set if calling from background thread.
        const IN_BACKGROUND = 0x1000_0000;
    }
}

bitflags! {
    /// Flag options to close files.
    #[derive(Default)]
    pub struct CloseFlags: u32 {
        /// No close option
        const NONE = 0x0000;
        /// Write into a temporary file and swap for crash safety.
        const SAFE_UPDATE = 0x0001;
    }
}

bitflags! {
    /// Result flag for file / format infos.
    #[derive(Default)]
    pub struct FormatOptionFlags: u32 {
        const NONE = 0;
        /// Can inject first-time XMP into an existing file.
        const CAN_INJECT_XMP = 0x0000_0001;
        /// Can expand XMP or other metadata in an existing file.
        const CAN_EXPAND = 0x0000_0002;
        /// Can copy one file to another, writing new metadata.
        const CAN_REWRITE = 0x0000_0004;
        /// Can expand, but prefers in-place update.
        const PREFERS_IN_PLACE = 0x0000_0008;
        /// Supports reconciliation between XMP and other forms.
        const CAN_RECONCILE = 0x0000_0010;
        /// Allows access to just the XMP, ignoring other forms.
        const ALLOWS_ONLY_XMP = 0x0000_0020;
        /// File handler returns raw XMP packet information.
        const RETURNS_RAW_PACKET = 0x0000_0040;
        /// The file handler does the file open and close.
        const HANDLER_OWNS_FILE = 0x0000_0100;
        /// The file handler allows crash-safe file updates.
        const ALLOW_SAFE_UPDATE = 0x0000_0200;
        /// The file format needs the XMP packet to be read-only.
        const NEEDS_READONLY_PACKET = 0x0000_0400;
        /// The file handler uses a "sidecar" file for the XMP.
        const USE_SIDECAR_XMP = 0x0000_0800;
        /// The format is folder oriented, for example the P2 video format.
        const FOLDER_BASED_FORMAT = 0x0000_1000;
    }
}

/// Wrap an XmpPacketInfo
pub struct PacketInfo(pub XmpPacketInfo);

impl Default for PacketInfo {
    fn default() -> PacketInfo {
        PacketInfo(XmpPacketInfo {
            offset: 0,
            length: 0,
            pad_size: 0,
            char_form: 0,
            writeable: false,
            has_wrapper: false,
            pad: 0,
        })
    }
}

pub struct XmpFile(*mut c::XmpFile);

impl Default for XmpFile {
    fn default() -> XmpFile {
        super::init();

        XmpFile(unsafe { c::xmp_files_new() })
    }
}

impl XmpFile {
    /// Create new XmpFile
    pub fn new() -> XmpFile {
        XmpFile::default()
    }

    /// Create and open a new XmpFile
    /// Equivalent to calling new then open.
    /// Return Err in case of failure
    pub fn new_from_file<P: AsRef<std::ffi::OsStr>>(
        path: P,
        options: OpenFlags,
    ) -> Result<XmpFile> {
        super::init();

        let path = path.as_ref().as_bytes();
        let pp = CString::new(path).unwrap();
        let ptr = unsafe { c::xmp_files_open_new(pp.as_ptr(), options.bits()) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(XmpFile(ptr))
    }

    /// Open an XmpFile. Usually called after new.
    pub fn open<P: AsRef<std::ffi::OsStr>>(&mut self, path: P, options: OpenFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        let path = path.as_ref().as_bytes();
        let pp = CString::new(path).unwrap();
        if unsafe { c::xmp_files_open(self.0, pp.as_ptr(), options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Close the XmpFile
    pub fn close(&mut self, options: CloseFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_close(self.0, options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if native pointer is null
    pub fn is_null(&self) -> bool {
        self.0.is_null()
    }

    /// Get a new Xmp fronm the currently open file
    pub fn get_new_xmp(&self) -> Result<Xmp> {
        let ptr = unsafe { c::xmp_files_get_new_xmp(self.0) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(Xmp::from(ptr))
    }

    /// Get the xmp data into an existing Xmp.
    pub fn get_xmp(&self, xmp: &mut Xmp) -> Result<()> {
        if self.is_null() || xmp.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_get_xmp(self.0, xmp.as_mut_ptr()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Get the xmp packet as a string.
    pub fn get_xmp_xmpstring(&self) -> Result<(XmpString, PacketInfo)> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        let mut packet = XmpString::new();
        let mut info = PacketInfo::default();
        if unsafe {
            c::xmp_files_get_xmp_xmpstring(
                self.0,
                packet.as_mut_ptr(),
                &mut info.0 as *mut XmpPacketInfo,
            )
        } {
            Ok((packet, info))
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if it can put the Xmp into the XmpFile.
    pub fn can_put_xmp(&self, xmp: &Xmp) -> bool {
        if self.is_null() || xmp.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp(self.0, xmp.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_xmpstring(&self, xmp_packet: &XmpString) -> bool {
        if self.is_null() || xmp_packet.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp_xmpstring(self.0, xmp_packet.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_str(&self, xmp_packet: &str) -> bool {
        if self.is_null() {
            return false;
        }
        let pp = CString::new(xmp_packet).unwrap();
        unsafe { c::xmp_files_can_put_xmp_cstr(self.0, pp.as_ptr(), xmp_packet.len()) }
    }

    /// Put the Xmp into the XmpFile
    pub fn put_xmp(&mut self, xmp: &Xmp) -> Result<()> {
        if self.is_null() || xmp.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_put_xmp(self.0, xmp.as_ptr()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Get info from the XmpFile.
    pub fn get_file_info(
        &self,
        file_path: &mut String,
        options: &mut OpenFlags,
        format: &mut FileType,
        handler_flags: &mut FormatOptionFlags,
    ) -> bool {
        if self.is_null() {
            return false;
        }
        let mut s = XmpString::new();

        let mut raw_options: u32 = 0;
        let mut raw_handler_flags: u32 = 0;
        let result = unsafe {
            c::xmp_files_get_file_info(
                self.0,
                s.as_mut_ptr(),
                &mut raw_options,
                format,
                &mut raw_handler_flags,
            )
        };
        *options = OpenFlags::from_bits(raw_options).unwrap_or_default();
        *handler_flags = FormatOptionFlags::from_bits(raw_handler_flags).unwrap_or_default();

        file_path.push_str(&s.to_string());
        result
    }

    /// Check the file format for the specified path
    pub fn check_file_format<P: AsRef<std::ffi::OsStr>>(path: P) -> FileType {
        let path = path.as_ref().as_bytes();
        let pp = CString::new(path).unwrap();
        unsafe { c::xmp_files_check_file_format(pp.as_ptr()) }
    }

    /// Get FormatOptions for the FileType
    pub fn get_format_info(format: FileType) -> Result<FormatOptionFlags> {
        let mut raw_options: u32 = 0;
        if unsafe { c::xmp_files_get_format_info(format, &mut raw_options) } {
            Ok(FormatOptionFlags::from_bits(raw_options).unwrap_or_default())
        } else {
            Err(crate::get_error())
        }
    }
}

impl Drop for XmpFile {
    /// Drop the XmpFile.
    fn drop(&mut self) {
        if !self.is_null() {
            unsafe { c::xmp_files_free(self.0) };
        }
    }
}

#[test]
fn it_works() {
    use crate::error::Error;

    let inited = crate::init();

    assert!(inited);

    let xf = XmpFile::new();
    assert!(!xf.is_null());

    assert!(crate::get_error() == Error::from(c::XmpError::Unknown));
}
<|endoftext|>


<commit_before>//! A simple application that downloads the data dictionary and creates code or
//! data to reproduce it in the core library.
//!
//! ### How to use
//!
//! Simply run the application. It will automatically retrieve the dictionary
//! from the DCMTK data dictionary and store the result in "entries.rs".
//! Future versions will enable different kinds of outputs.
//!
//! Please use the `--help` flag for the full usage information.

use clap::{App, Arg};
use serde::Serialize;
use regex::Regex;

use heck::ShoutySnakeCase;
use std::fs::{create_dir_all, File};
use std::io::{BufRead, BufReader, Write};
use std::path::Path;

/// url to DCMTK dic file
const DEFAULT_LOCATION: &str = "https://raw.githubusercontent.com/DCMTK/dcmtk/master/dcmdata/data/dicom.dic";

fn main() {
    let matches = App::new("DICOM Dictionary Builder")
        .version("0.1.0")
        .arg(
            Arg::with_name("FROM")
                .default_value(DEFAULT_LOCATION)
                .help("Where to fetch the dictionary from"),
        )
        .arg(
            Arg::with_name("no-retired")
                .help("Whether to ignore retired tags")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("OUTPUT")
                .short("o")
                .help("The path to the output file")
                .default_value("tags.rs")
                .takes_value(true),
        )
        .get_matches();

    let ignore_retired = matches.is_present("no-retired");

    let src = matches.value_of("FROM").unwrap();

    let dst = Path::new(matches.value_of("OUTPUT").unwrap());

    if src.starts_with("http:") || src.starts_with("https:") {
        // read from URL
        //let src = Uri::from_str(src).unwrap();
        println!("Downloading DICOM dictionary ...");
        let mut resp = reqwest::get(src).unwrap();
        let mut data = vec![];
        resp.copy_to(&mut data).unwrap();
        let entries = parse_entries(&*data).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired).expect("Failed to write file");
    } else {
        // read from File
        let file = File::open(src).unwrap();
        let entries = parse_entries(BufReader::new(file)).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired).expect("Failed to write file");
    }
}

type DynResult<T> = Result<T, Box<dyn std::error::Error>>;

fn parse_entries<R: BufRead>(source: R) -> DynResult<Vec<Entry>> {
    let mut result = vec![];

    for line in source.lines() {
        let line = line?;
        if line.starts_with('#') {
            continue;
        }

        // (0010,0010)	PN	PatientName	1	DICOM
        let parts: Vec<&str> = line.split('\t').collect();

        if parts[4] == "ILLEGAL" || parts[4] == "PRIVATE" || parts[4] == "GENERIC" {
            continue;
        }

        let mut vr = parts[1].to_string();

        // Some fields are dependent upon context.
        // We may want to support this at some point, but for now,
        // let's just choose some good defaults.
        if vr == "up" {
            vr = "UL".to_string();
        }
        if vr == "xs" {
            vr = "US".to_string();
        }
        if vr == "ox" {
            vr = "OB".to_string();
        }
        if vr == "lt" {
            vr = "OW".to_string();
        }
        if vr == "na" {
            // These are the "Item", "ItemDelimitationItem", "SequenceDelimitationItem", etc, values.
            // Question, should we generate const values for these, and use them internally?
            continue;
        }

        let mut alias = parts[2].to_string();
        if alias.starts_with("RETIRED_") {
            alias = alias.trim_start_matches("RETIRED_").to_string();
        }

        let tag = parts[0].to_string();

        let regex_tag = Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{4})\)$")?;
        let regex_tag_group100 = Regex::new(r"^\(([0-9A-F]{2})00-[0-9A-F]{2}FF,([0-9A-F]{4})\)$")?;
        let regex_tag_element100 =
            Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{2})00-[0-9A-F]{2}FF\)$")?;

        let cap = regex_tag.captures(tag.as_str());
        let tag_type;
        let tag_declaration = if let Some(cap) = cap {
            // single tag
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Single;
            format!("Tag(0x{}, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_group100.captures(tag.as_str()) {
            // tag range over groups: (ggxx, eeee)
            let group = cap.get(1).expect("capture group 1: group portion").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Group100;
            format!("Tag(0x{}00, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_element100.captures(tag.as_str()) {
            // tag range over elements: (gggg, eexx)
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap
                .get(2)
                .expect("capture group 2: element portion")
                .as_str();
            tag_type = TagType::Element100;
            format!("Tag(0x{}, 0x{}00)", group, elem)
        } else {
            panic!("invalid tag: {}", alias);
        };

        result.push(Entry {
            tag,
            vr,
            alias,
            vm: parts[3].to_string(),
            obs: parts[4].to_string(),
            is_retired: parts[4].contains("retired"),
            tag_declaration,
            tag_type,
        });
    }

    Ok(result)
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
enum TagType {
    Single,
    Group100,
    Element100,
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
struct Entry {
    /// (0010,0010)
    tag: String,
    /// PN
    vr: String,
    /// PatientName
    alias: String,
    /// 1
    vm: String,
    /// DICOM
    obs: String,
    /// Retired field?
    is_retired: bool,
    /// Tag declaration "Tag(0x6000, 0x1102)"
    tag_declaration: String,
    /// The type the tag represents
    tag_type: TagType,
}

fn to_code_file<P: AsRef<Path>>(
    dest_path: P,
    entries: Vec<Entry>,
    include_retired: bool,
) -> DynResult<()> {
    if let Some(p_dir) = dest_path.as_ref().parent() {
        create_dir_all(&p_dir)?;
    }
    let mut f = File::create(&dest_path)?;

    f.write_all(
        b"//! Automatically generated. Edit at your own risk.\n\n\
    use dicom_core::dictionary::{DictionaryEntryRef, TagRange, TagRange::*};\n\
    use dicom_core::Tag;\n\
    use dicom_core::VR::*;\n\n",
    )?;

    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        writeln!(f, "/// {} {} {} {} {}", e.alias, e.tag, e.vr, e.vm, e.obs)?;

        let tag_type = match e.tag_type {
            TagType::Single => "Tag",
            _ => "TagRange",
        };

        let tag_set = match e.tag_type {
            TagType::Single => e.tag_declaration.clone(),
            TagType::Group100 => format!("Group100({})", e.tag_declaration),
            TagType::Element100 => format!("Element100({})", e.tag_declaration),
        };

        writeln!(
            f,
            "#[rustfmt::skip]\npub const {}: {} = {};",
            e.alias.to_shouty_snake_case(),
            tag_type,
            tag_set,
        )?;
    }

    f.write_all(
        b"\n\n\
    type E = DictionaryEntryRef<'static>;\n\n\
    #[rustfmt::skip]\n\
    pub(crate) const ENTRIES: &[E] = &[\n",
    )?;
    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        let (vr1, vr2) = e.vr.split_at(2);

        let second_vr = if vr2 != "" {
            format!(" /*{} */", vr2)
        } else {
            vr2.to_string()
        };

        let tag_set = match e.tag_type {
            TagType::Single => format!("Single({})", e.alias.to_shouty_snake_case()),
            _ => e.alias.to_shouty_snake_case(),
        };

        writeln!(
            f,
            "    E {{ tag: {}, alias: \"{}\", vr: {}{} }}, // {}",
            tag_set, e.alias, vr1, second_vr, e.obs
        )?;
    }
    f.write_all(b"];\n")?;

    Ok(())
}
<commit_msg>[dictionary-builder] add context-dependent VR "px"<commit_after>//! A simple application that downloads the data dictionary and creates code or
//! data to reproduce it in the core library.
//!
//! ### How to use
//!
//! Simply run the application. It will automatically retrieve the dictionary
//! from the DCMTK data dictionary and store the result in "entries.rs".
//! Future versions will enable different kinds of outputs.
//!
//! Please use the `--help` flag for the full usage information.

use clap::{App, Arg};
use serde::Serialize;
use regex::Regex;

use heck::ShoutySnakeCase;
use std::fs::{create_dir_all, File};
use std::io::{BufRead, BufReader, Write};
use std::path::Path;

/// url to DCMTK dic file
const DEFAULT_LOCATION: &str = "https://raw.githubusercontent.com/DCMTK/dcmtk/master/dcmdata/data/dicom.dic";

fn main() {
    let matches = App::new("DICOM Dictionary Builder")
        .version("0.1.0")
        .arg(
            Arg::with_name("FROM")
                .default_value(DEFAULT_LOCATION)
                .help("Where to fetch the dictionary from"),
        )
        .arg(
            Arg::with_name("no-retired")
                .help("Whether to ignore retired tags")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("OUTPUT")
                .short("o")
                .help("The path to the output file")
                .default_value("tags.rs")
                .takes_value(true),
        )
        .get_matches();

    let ignore_retired = matches.is_present("no-retired");

    let src = matches.value_of("FROM").unwrap();

    let dst = Path::new(matches.value_of("OUTPUT").unwrap());

    if src.starts_with("http:") || src.starts_with("https:") {
        // read from URL
        //let src = Uri::from_str(src).unwrap();
        println!("Downloading DICOM dictionary ...");
        let mut resp = reqwest::get(src).unwrap();
        let mut data = vec![];
        resp.copy_to(&mut data).unwrap();
        let entries = parse_entries(&*data).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired).expect("Failed to write file");
    } else {
        // read from File
        let file = File::open(src).unwrap();
        let entries = parse_entries(BufReader::new(file)).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired).expect("Failed to write file");
    }
}

type DynResult<T> = Result<T, Box<dyn std::error::Error>>;

fn parse_entries<R: BufRead>(source: R) -> DynResult<Vec<Entry>> {
    let mut result = vec![];

    for line in source.lines() {
        let line = line?;
        if line.starts_with('#') {
            continue;
        }

        // (0010,0010)	PN	PatientName	1	DICOM
        let parts: Vec<&str> = line.split('\t').collect();

        if parts[4] == "ILLEGAL" || parts[4] == "PRIVATE" || parts[4] == "GENERIC" {
            continue;
        }

        let mut vr = parts[1].to_string();

        // Some fields are dependent upon context.
        // We may want to support this at some point, but for now,
        // let's just choose some good defaults.
        if vr == "up" {
            vr = "UL".to_string();
        }
        if vr == "xs" {
            vr = "US".to_string();
        }
        if vr == "ox" {
            vr = "OB".to_string();
        }
        if vr == "px" {
            vr = "OB".to_string();
        }
        if vr == "lt" {
            vr = "OW".to_string();
        }
        if vr == "na" {
            // These are the "Item", "ItemDelimitationItem", "SequenceDelimitationItem", etc, values.
            // Question, should we generate const values for these, and use them internally?
            continue;
        }

        let mut alias = parts[2].to_string();
        if alias.starts_with("RETIRED_") {
            alias = alias.trim_start_matches("RETIRED_").to_string();
        }

        let tag = parts[0].to_string();

        let regex_tag = Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{4})\)$")?;
        let regex_tag_group100 = Regex::new(r"^\(([0-9A-F]{2})00-[0-9A-F]{2}FF,([0-9A-F]{4})\)$")?;
        let regex_tag_element100 =
            Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{2})00-[0-9A-F]{2}FF\)$")?;

        let cap = regex_tag.captures(tag.as_str());
        let tag_type;
        let tag_declaration = if let Some(cap) = cap {
            // single tag
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Single;
            format!("Tag(0x{}, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_group100.captures(tag.as_str()) {
            // tag range over groups: (ggxx, eeee)
            let group = cap.get(1).expect("capture group 1: group portion").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Group100;
            format!("Tag(0x{}00, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_element100.captures(tag.as_str()) {
            // tag range over elements: (gggg, eexx)
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap
                .get(2)
                .expect("capture group 2: element portion")
                .as_str();
            tag_type = TagType::Element100;
            format!("Tag(0x{}, 0x{}00)", group, elem)
        } else {
            panic!("invalid tag: {}", alias);
        };

        result.push(Entry {
            tag,
            vr,
            alias,
            vm: parts[3].to_string(),
            obs: parts[4].to_string(),
            is_retired: parts[4].contains("retired"),
            tag_declaration,
            tag_type,
        });
    }

    Ok(result)
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
enum TagType {
    Single,
    Group100,
    Element100,
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
struct Entry {
    /// (0010,0010)
    tag: String,
    /// PN
    vr: String,
    /// PatientName
    alias: String,
    /// 1
    vm: String,
    /// DICOM
    obs: String,
    /// Retired field?
    is_retired: bool,
    /// Tag declaration "Tag(0x6000, 0x1102)"
    tag_declaration: String,
    /// The type the tag represents
    tag_type: TagType,
}

fn to_code_file<P: AsRef<Path>>(
    dest_path: P,
    entries: Vec<Entry>,
    include_retired: bool,
) -> DynResult<()> {
    if let Some(p_dir) = dest_path.as_ref().parent() {
        create_dir_all(&p_dir)?;
    }
    let mut f = File::create(&dest_path)?;

    f.write_all(
        b"//! Automatically generated. Edit at your own risk.\n\n\
    use dicom_core::dictionary::{DictionaryEntryRef, TagRange, TagRange::*};\n\
    use dicom_core::Tag;\n\
    use dicom_core::VR::*;\n\n",
    )?;

    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        writeln!(f, "/// {} {} {} {} {}", e.alias, e.tag, e.vr, e.vm, e.obs)?;

        let tag_type = match e.tag_type {
            TagType::Single => "Tag",
            _ => "TagRange",
        };

        let tag_set = match e.tag_type {
            TagType::Single => e.tag_declaration.clone(),
            TagType::Group100 => format!("Group100({})", e.tag_declaration),
            TagType::Element100 => format!("Element100({})", e.tag_declaration),
        };

        writeln!(
            f,
            "#[rustfmt::skip]\npub const {}: {} = {};",
            e.alias.to_shouty_snake_case(),
            tag_type,
            tag_set,
        )?;
    }

    f.write_all(
        b"\n\n\
    type E = DictionaryEntryRef<'static>;\n\n\
    #[rustfmt::skip]\n\
    pub(crate) const ENTRIES: &[E] = &[\n",
    )?;
    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        let (vr1, vr2) = e.vr.split_at(2);

        let second_vr = if vr2 != "" {
            format!(" /*{} */", vr2)
        } else {
            vr2.to_string()
        };

        let tag_set = match e.tag_type {
            TagType::Single => format!("Single({})", e.alias.to_shouty_snake_case()),
            _ => e.alias.to_shouty_snake_case(),
        };

        writeln!(
            f,
            "    E {{ tag: {}, alias: \"{}\", vr: {}{} }}, // {}",
            tag_set, e.alias, vr1, second_vr, e.obs
        )?;
    }
    f.write_all(b"];\n")?;

    Ok(())
}
<|endoftext|>


<commit_before>use columnvalueops::ColumnValueOps;
use databaseinfo::{DatabaseInfo, TableInfo};
use identifier::Identifier;
use sqlsyntax::ast;

use std::fmt;

mod execute;
mod sexpression;
mod source;
pub use self::execute::*;
pub use self::sexpression::*;
use self::source::*;

pub enum QueryPlanCompileError {
    TableDoesNotExist(Identifier),
    /// ambiguous column name; two or more tables have a column of the same name
    AmbiguousColumnName(Identifier),
    BadIdentifier(String),
    BadStringLiteral(String),
    BadNumberLiteral(String),
    UnknownFunctionName(Identifier),
    AggregateFunctionRequiresOneArgument
}

impl fmt::Display for QueryPlanCompileError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        use self::QueryPlanCompileError::*;

        match self {
            &TableDoesNotExist(ref name) => {
                write!(f, "table does not exist: {}", name)
            },
            &AmbiguousColumnName(ref name) => {
                write!(f, "ambiguous column name: {}", name)
            },
            &BadIdentifier(ref name) => {
                write!(f, "bad identifier: {}", name)
            },
            &BadStringLiteral(ref s) => {
                write!(f, "bad string literal: {}", s)
            },
            &BadNumberLiteral(ref s) => {
                write!(f, "bad number literal: {}", s)
            },
            &UnknownFunctionName(ref s) => {
                write!(f, "unknown function name: {}", s)
            },
            &AggregateFunctionRequiresOneArgument => {
                write!(f, "aggregate function requires exactly one argument")
            },
        }
    }
}

pub struct QueryPlan<'a, DB: DatabaseInfo>
where <DB as DatabaseInfo>::Table: 'a
{
    pub expr: SExpression<'a, DB>,
    pub out_column_names: Vec<Identifier>
}

fn new_identifier(value: &str) -> Result<Identifier, QueryPlanCompileError> {
    Identifier::new(value).ok_or(QueryPlanCompileError::BadIdentifier(value.to_string()))
}

impl<'a, DB: DatabaseInfo> QueryPlan<'a, DB>
where <DB as DatabaseInfo>::Table: 'a
{
    pub fn compile_select(db: &'a DB, stmt: ast::SelectStatement)
    -> Result<QueryPlan<'a, DB>, QueryPlanCompileError>
    {
        let scope = SourceScope::new(None, Vec::new(), Vec::new());

        let mut source_id: u32 = 0;

        let source_id_fn = || {
            let old_source_id = source_id;
            source_id += 1;
            old_source_id
        };

        let mut compiler = Compiler {
            db: db,
            new_source_id: source_id_fn
        };

        compiler.compile_select_recurse(stmt, &scope)
    }
}

struct Compiler<'a, DB: DatabaseInfo, F>
where DB: 'a, <DB as DatabaseInfo>::Table: 'a, F: FnMut() -> u32
{
    db: &'a DB,
    new_source_id: F
}

struct FromWhere<'a, DB: DatabaseInfo>
where <DB as DatabaseInfo>::Table: 'a
{
    tables: Vec<FromWhereTableOrSubquery<'a, DB>>,
    where_expr: Option<SExpression<'a, DB>>
}

enum FromWhereTableOrSubquery<'a, DB: DatabaseInfo>
where <DB as DatabaseInfo>::Table: 'a
{
    Table {
        source_id: u32,
        table: &'a <DB as DatabaseInfo>::Table
    },
    Subquery {
        source_id: u32,
        expr: SExpression<'a, DB>
    }
}

impl<'a, DB: DatabaseInfo> FromWhere<'a, DB>
where <DB as DatabaseInfo>::Table: 'a
{
    fn evaluate(self, inner_expr: SExpression<'a, DB>) -> SExpression<'a, DB> {
        let core_expr = if let Some(where_expr) = self.where_expr {
            SExpression::If {
                predicate: Box::new(where_expr),
                yield_fn: Box::new(inner_expr)
            }
        } else {
            inner_expr
        };

        self.tables.into_iter().fold(core_expr, |nested_expr, x| {
            match x {
                FromWhereTableOrSubquery::Subquery { source_id, expr } => {
                    SExpression::Map {
                        source_id: source_id,
                        yield_in_fn: Box::new(expr),
                        yield_out_fn: Box::new(nested_expr)
                    }
                },
                FromWhereTableOrSubquery::Table { source_id, table } => {
                    SExpression::Scan {
                        source_id: source_id,
                        table: table,
                        yield_fn: Box::new(nested_expr)
                    }
                }
            }
        })
    }
}

impl<'a, DB: DatabaseInfo, F> Compiler<'a, DB, F>
where DB: 'a, <DB as DatabaseInfo>::Table: 'a, F: FnMut() -> u32
{
    fn new_source_id(&mut self) -> u32 {
        (&mut self.new_source_id)()
    }

    fn compile_select_recurse<'b>(&mut self, stmt: ast::SelectStatement, outer_scope: &'b SourceScope<'b>)
    -> Result<QueryPlan<'a, DB>, QueryPlanCompileError>
    {
        // Unimplemented syntaxes: GROUP BY, HAVING, ORDER BY
        // TODO - implement them!
        if !stmt.group_by.is_empty() { unimplemented!() }
        if stmt.having.is_some() { unimplemented!() }
        if !stmt.order_by.is_empty() { unimplemented!() }

        let mut arbitrary_column_count = 0;

        let mut arbitrary_column_name = || {
            let s = format!("_{}", arbitrary_column_count);
            arbitrary_column_count += 1;

            Identifier::new(&s).unwrap()
        };

        // FROM and WHERE are compiled together.
        // This makes sense for INNER and OUTER joins, which also
        // contain ON (conditional) expressions.

        let (new_scope, from_where) = try!(self.from_where(stmt.from, stmt.where_expr, outer_scope));

        let (column_names, select_exprs) = {
            let mut a: Vec<_> = Vec::new();
            for c in stmt.result_columns {
                match c {
                    ast::SelectColumn::AllColumns => {
                        a.extend(new_scope.tables().iter().flat_map(|table| {
                            let source_id = table.source_id;

                            table.out_column_names.iter().enumerate().map(move |(i, name)| {
                                (name.clone(), SExpression::ColumnField {
                                    source_id: source_id,
                                    column_offset: i as u32
                                })
                            })
                        }));
                    },
                    ast::SelectColumn::Expr { expr, alias } => {
                        let column_name = if let Some(alias) = alias {
                            try!(new_identifier(&alias))
                        } else {
                            // if the expression is a simple identifier, make that
                            // the column name. else, assign an arbitrary name.
                            if let &ast::Expression::Ident(ref n) = &expr {
                                try!(new_identifier(n))
                            } else {
                                arbitrary_column_name()
                            }
                        };

                        let e = try!(self.ast_expression_to_sexpression(expr, &new_scope));
                        a.push((column_name, e));
                    }
                }
            }

            a.into_iter().unzip()
        };

        let expr = from_where.evaluate(SExpression::Yield { fields: select_exprs });

        Ok(QueryPlan {
            expr: expr,
            out_column_names: column_names
        })
    }

    fn from_where<'b>(&mut self, from: ast::From, where_expr: Option<ast::Expression>, scope: &'b SourceScope<'b>)
    -> Result<(SourceScope<'b>, FromWhere<'a, DB>), QueryPlanCompileError>
    {
        // TODO - avoid naive nested scans when indices are available

        // All FROM subqueries are nested, never correlated.
        let ast_cross_tables = match from {
            ast::From::Cross(v) => v,
            ast::From::Join {..} => unimplemented!()
        };

        let a: Vec<_> = try!(ast_cross_tables.into_iter().map(|ast_table_or_subquery| {
            match ast_table_or_subquery {
                ast::TableOrSubquery::Subquery { subquery, alias } => {
                    let plan = try!(self.compile_select_recurse(*subquery, scope));
                    let alias_identifier = try!(new_identifier(&alias));

                    let source_id = self.new_source_id();

                    let s = TableOrSubquery {
                        source_id: source_id,
                        out_column_names: plan.out_column_names
                    };

                    let t = FromWhereTableOrSubquery::Subquery {
                        source_id: source_id,
                        expr: plan.expr
                    };

                    Ok(((s, t), alias_identifier))
                },
                ast::TableOrSubquery::Table { table, alias } => {
                    let table_name_identifier = try!(new_identifier(&table.table_name));
                    let table = match self.db.find_table_by_name(&table_name_identifier) {
                        Some(table) => table,
                        None => return Err(QueryPlanCompileError::TableDoesNotExist(table_name_identifier))
                    };

                    let alias_identifier = if let Some(alias) = alias {
                        try!(new_identifier(&alias))
                    } else {
                        table_name_identifier
                    };

                    let source_id = self.new_source_id();

                    let s = TableOrSubquery {
                        source_id: source_id,
                        out_column_names: table.get_column_names()
                    };

                    let t = FromWhereTableOrSubquery::Table {
                        source_id: source_id,
                        table: table
                    };

                    Ok(((s, t), alias_identifier))
                }
            }
        }).collect());

        let (tables, table_aliases): (Vec<_>, _) = a.into_iter().unzip();

        let (source_tables, fromwhere_tables) = tables.into_iter().unzip();

        let new_scope = SourceScope::new(Some(scope), source_tables, table_aliases);

        let where_expr = if let Some(where_expr) = where_expr {
            Some(try!(self.ast_expression_to_sexpression(where_expr, &new_scope)))
        } else {
            None
        };

        Ok((new_scope, FromWhere {
            tables: fromwhere_tables,
            where_expr: where_expr
        }))
    }

    fn ast_expression_to_sexpression<'b>(&mut self, ast: ast::Expression, scope: &'b SourceScope<'b>)
    -> Result<SExpression<'a, DB>, QueryPlanCompileError>
    {
        use std::borrow::IntoCow;

        match ast {
            ast::Expression::Ident(s) => {
                let column_identifier = try!(new_identifier(&s));

                let (source_id, column_offset) = match scope.get_column_offset(&column_identifier) {
                    Some(v) => v,
                    None => return Err(QueryPlanCompileError::AmbiguousColumnName(column_identifier))
                };

                Ok(SExpression::ColumnField {
                    source_id: source_id,
                    column_offset: column_offset
                })
            },
            ast::Expression::IdentMember(s1, s2) => {
                let table_identifier = try!(new_identifier(&s1));
                let column_identifier = try!(new_identifier(&s2));

                let (source_id, column_offset) = match scope.get_table_column_offset(&table_identifier, &column_identifier) {
                    Some(v) => v,
                    None => return Err(QueryPlanCompileError::AmbiguousColumnName(column_identifier))
                };

                Ok(SExpression::ColumnField {
                    source_id: source_id,
                    column_offset: column_offset
                })
            },
            ast::Expression::BinaryOp { lhs, rhs, op } => {
                let l = try!(self.ast_expression_to_sexpression(*lhs, scope));
                let r = try!(self.ast_expression_to_sexpression(*rhs, scope));

                Ok(SExpression::BinaryOp {
                    op: ast_binaryop_to_sexpression_binaryop(op),
                    lhs: Box::new(l),
                    rhs: Box::new(r)
                })
            },
            ast::Expression::StringLiteral(s) => {
                match DB::ColumnValue::from_string_literal(s.into_cow()) {
                    Ok(value) => Ok(SExpression::Value(value)),
                    Err(s) => Err(QueryPlanCompileError::BadStringLiteral(s.into_owned()))
                }
            },
            ast::Expression::Number(s) => {
                match DB::ColumnValue::from_number_literal(s.into_cow()) {
                    Ok(value) => Ok(SExpression::Value(value)),
                    Err(s) => Err(QueryPlanCompileError::BadNumberLiteral(s.into_owned()))
                }
            },
            ast::Expression::Subquery(subquery) => {
                let source_id = self.new_source_id();

                let plan = try!(self.compile_select_recurse(*subquery, scope));

                Ok(SExpression::Map {
                    source_id: source_id,
                    yield_in_fn: Box::new(plan.expr),
                    yield_out_fn: Box::new(SExpression::ColumnField {
                        source_id: source_id,
                        column_offset: 0
                    })
                })
            },
            ast::Expression::FunctionCall { name, arguments } => {
                let ident = try!(new_identifier(&name));

                macro_rules! aggregate {
                    ($op:expr) => (
                        if arguments.len() != 1 {
                            Err(QueryPlanCompileError::AggregateFunctionRequiresOneArgument)
                        } else {
                            let arg = arguments.into_iter().nth(0).unwrap();
                            let value = try!(self.ast_expression_to_sexpression(arg, scope));
                            let source_id = unimplemented!();
                            Ok(SExpression::AggregateOp {
                                op: $op,
                                source_id: source_id,
                                value: Box::new(value)
                            })
                        }
                    )
                }

                match &ident as &str {
                    "count" => aggregate!(AggregateOp::Count),
                    "avg" => aggregate!(AggregateOp::Avg),
                    _ => Err(QueryPlanCompileError::UnknownFunctionName(ident))
                }
            },
            e => panic!("unimplemented: {:?}", e)
        }
    }
}

fn ast_binaryop_to_sexpression_binaryop(ast: ast::BinaryOp) -> BinaryOp {
    match ast {
        ast::BinaryOp::Equal => BinaryOp::Equal,
        ast::BinaryOp::NotEqual => BinaryOp::NotEqual,
        ast::BinaryOp::LessThan => BinaryOp::LessThan,
        ast::BinaryOp::LessThanOrEqual => BinaryOp::LessThanOrEqual,
        ast::BinaryOp::GreaterThan => BinaryOp::GreaterThan,
        ast::BinaryOp::GreaterThanOrEqual => BinaryOp::GreaterThanOrEqual,
        ast::BinaryOp::And => BinaryOp::And,
        ast::BinaryOp::Or => BinaryOp::Or,
        ast::BinaryOp::Add => BinaryOp::Add,
        ast::BinaryOp::Subtract => BinaryOp::Subtract,
        ast::BinaryOp::Multiply => BinaryOp::Multiply,
        ast::BinaryOp::BitAnd => BinaryOp::BitAnd,
        ast::BinaryOp::BitOr => BinaryOp::BitOr,
        ast::BinaryOp::Concatenate => BinaryOp::Concatenate,
    }
}

impl<'a, DB: DatabaseInfo> fmt::Display for QueryPlan<'a, DB>
where <DB as DatabaseInfo>::Table: 'a
{
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let cn: Vec<_> = self.out_column_names.iter().map(|n| format!("`{}`", n)).collect();

        try!(writeln!(f, "query plan"));
        try!(writeln!(f, "column names: ({})", cn.connect(", ")));
        self.expr.fmt(f)
    }
}
<commit_msg>Move SELECT compilation into its own method<commit_after>use columnvalueops::ColumnValueOps;
use databaseinfo::{DatabaseInfo, TableInfo};
use identifier::Identifier;
use sqlsyntax::ast;

use std::fmt;

mod execute;
mod sexpression;
mod source;
pub use self::execute::*;
pub use self::sexpression::*;
use self::source::*;

pub enum QueryPlanCompileError {
    TableDoesNotExist(Identifier),
    /// ambiguous column name; two or more tables have a column of the same name
    AmbiguousColumnName(Identifier),
    BadIdentifier(String),
    BadStringLiteral(String),
    BadNumberLiteral(String),
    UnknownFunctionName(Identifier),
    AggregateFunctionRequiresOneArgument
}

impl fmt::Display for QueryPlanCompileError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        use self::QueryPlanCompileError::*;

        match self {
            &TableDoesNotExist(ref name) => {
                write!(f, "table does not exist: {}", name)
            },
            &AmbiguousColumnName(ref name) => {
                write!(f, "ambiguous column name: {}", name)
            },
            &BadIdentifier(ref name) => {
                write!(f, "bad identifier: {}", name)
            },
            &BadStringLiteral(ref s) => {
                write!(f, "bad string literal: {}", s)
            },
            &BadNumberLiteral(ref s) => {
                write!(f, "bad number literal: {}", s)
            },
            &UnknownFunctionName(ref s) => {
                write!(f, "unknown function name: {}", s)
            },
            &AggregateFunctionRequiresOneArgument => {
                write!(f, "aggregate function requires exactly one argument")
            },
        }
    }
}

pub struct QueryPlan<'a, DB: DatabaseInfo>
where <DB as DatabaseInfo>::Table: 'a
{
    pub expr: SExpression<'a, DB>,
    pub out_column_names: Vec<Identifier>
}

fn new_identifier(value: &str) -> Result<Identifier, QueryPlanCompileError> {
    Identifier::new(value).ok_or(QueryPlanCompileError::BadIdentifier(value.to_string()))
}

impl<'a, DB: DatabaseInfo> QueryPlan<'a, DB>
where <DB as DatabaseInfo>::Table: 'a
{
    pub fn compile_select(db: &'a DB, stmt: ast::SelectStatement)
    -> Result<QueryPlan<'a, DB>, QueryPlanCompileError>
    {
        let scope = SourceScope::new(None, Vec::new(), Vec::new());

        let mut source_id: u32 = 0;

        let source_id_fn = || {
            let old_source_id = source_id;
            source_id += 1;
            old_source_id
        };

        let mut compiler = Compiler {
            db: db,
            new_source_id: source_id_fn
        };

        compiler.compile_select_recurse(stmt, &scope)
    }
}

struct Compiler<'a, DB: DatabaseInfo, F>
where DB: 'a, <DB as DatabaseInfo>::Table: 'a, F: FnMut() -> u32
{
    db: &'a DB,
    new_source_id: F
}

struct FromWhere<'a, DB: DatabaseInfo>
where <DB as DatabaseInfo>::Table: 'a
{
    tables: Vec<FromWhereTableOrSubquery<'a, DB>>,
    where_expr: Option<SExpression<'a, DB>>
}

enum FromWhereTableOrSubquery<'a, DB: DatabaseInfo>
where <DB as DatabaseInfo>::Table: 'a
{
    Table {
        source_id: u32,
        table: &'a <DB as DatabaseInfo>::Table
    },
    Subquery {
        source_id: u32,
        expr: SExpression<'a, DB>
    }
}

impl<'a, DB: DatabaseInfo> FromWhere<'a, DB>
where <DB as DatabaseInfo>::Table: 'a
{
    fn evaluate(self, inner_expr: SExpression<'a, DB>) -> SExpression<'a, DB> {
        let core_expr = if let Some(where_expr) = self.where_expr {
            SExpression::If {
                predicate: Box::new(where_expr),
                yield_fn: Box::new(inner_expr)
            }
        } else {
            inner_expr
        };

        self.tables.into_iter().fold(core_expr, |nested_expr, x| {
            match x {
                FromWhereTableOrSubquery::Subquery { source_id, expr } => {
                    SExpression::Map {
                        source_id: source_id,
                        yield_in_fn: Box::new(expr),
                        yield_out_fn: Box::new(nested_expr)
                    }
                },
                FromWhereTableOrSubquery::Table { source_id, table } => {
                    SExpression::Scan {
                        source_id: source_id,
                        table: table,
                        yield_fn: Box::new(nested_expr)
                    }
                }
            }
        })
    }
}

impl<'a, DB: DatabaseInfo, F> Compiler<'a, DB, F>
where DB: 'a, <DB as DatabaseInfo>::Table: 'a, F: FnMut() -> u32
{
    fn new_source_id(&mut self) -> u32 {
        (&mut self.new_source_id)()
    }

    fn compile_select_recurse<'b>(&mut self, stmt: ast::SelectStatement, outer_scope: &'b SourceScope<'b>)
    -> Result<QueryPlan<'a, DB>, QueryPlanCompileError>
    {
        // Unimplemented syntaxes: GROUP BY, HAVING, ORDER BY
        // TODO - implement them!
        if !stmt.group_by.is_empty() { unimplemented!() }
        if stmt.having.is_some() { unimplemented!() }
        if !stmt.order_by.is_empty() { unimplemented!() }

        // FROM and WHERE are compiled together.
        // This makes sense for INNER and OUTER joins, which also
        // contain ON (conditional) expressions.

        let (new_scope, from_where) = try!(self.from_where(stmt.from, stmt.where_expr, outer_scope));

        let (column_names, select_exprs) = try!(self.select(stmt.result_columns, &new_scope));

        let expr = from_where.evaluate(SExpression::Yield { fields: select_exprs });

        Ok(QueryPlan {
            expr: expr,
            out_column_names: column_names
        })
    }

    fn from_where<'b>(&mut self, from: ast::From, where_expr: Option<ast::Expression>, scope: &'b SourceScope<'b>)
    -> Result<(SourceScope<'b>, FromWhere<'a, DB>), QueryPlanCompileError>
    {
        // TODO - avoid naive nested scans when indices are available

        // All FROM subqueries are nested, never correlated.
        let ast_cross_tables = match from {
            ast::From::Cross(v) => v,
            ast::From::Join {..} => unimplemented!()
        };

        let a: Vec<_> = try!(ast_cross_tables.into_iter().map(|ast_table_or_subquery| {
            match ast_table_or_subquery {
                ast::TableOrSubquery::Subquery { subquery, alias } => {
                    let plan = try!(self.compile_select_recurse(*subquery, scope));
                    let alias_identifier = try!(new_identifier(&alias));

                    let source_id = self.new_source_id();

                    let s = TableOrSubquery {
                        source_id: source_id,
                        out_column_names: plan.out_column_names
                    };

                    let t = FromWhereTableOrSubquery::Subquery {
                        source_id: source_id,
                        expr: plan.expr
                    };

                    Ok(((s, t), alias_identifier))
                },
                ast::TableOrSubquery::Table { table, alias } => {
                    let table_name_identifier = try!(new_identifier(&table.table_name));
                    let table = match self.db.find_table_by_name(&table_name_identifier) {
                        Some(table) => table,
                        None => return Err(QueryPlanCompileError::TableDoesNotExist(table_name_identifier))
                    };

                    let alias_identifier = if let Some(alias) = alias {
                        try!(new_identifier(&alias))
                    } else {
                        table_name_identifier
                    };

                    let source_id = self.new_source_id();

                    let s = TableOrSubquery {
                        source_id: source_id,
                        out_column_names: table.get_column_names()
                    };

                    let t = FromWhereTableOrSubquery::Table {
                        source_id: source_id,
                        table: table
                    };

                    Ok(((s, t), alias_identifier))
                }
            }
        }).collect());

        let (tables, table_aliases): (Vec<_>, _) = a.into_iter().unzip();

        let (source_tables, fromwhere_tables) = tables.into_iter().unzip();

        let new_scope = SourceScope::new(Some(scope), source_tables, table_aliases);

        let where_expr = if let Some(where_expr) = where_expr {
            Some(try!(self.ast_expression_to_sexpression(where_expr, &new_scope)))
        } else {
            None
        };

        Ok((new_scope, FromWhere {
            tables: fromwhere_tables,
            where_expr: where_expr
        }))
    }

    fn select<'b>(&mut self, result_columns: Vec<ast::SelectColumn>, scope: &'b SourceScope<'b>)
    -> Result<(Vec<Identifier>, Vec<SExpression<'a, DB>>), QueryPlanCompileError>
    {
        let mut arbitrary_column_count = 0;

        let mut arbitrary_column_name = || {
            let s = format!("_{}", arbitrary_column_count);
            arbitrary_column_count += 1;

            Identifier::new(&s).unwrap()
        };

        let mut a: Vec<_> = Vec::new();

        for c in result_columns {
            match c {
                ast::SelectColumn::AllColumns => {
                    a.extend(scope.tables().iter().flat_map(|table| {
                        let source_id = table.source_id;

                        table.out_column_names.iter().enumerate().map(move |(i, name)| {
                            (name.clone(), SExpression::ColumnField {
                                source_id: source_id,
                                column_offset: i as u32
                            })
                        })
                    }));
                },
                ast::SelectColumn::Expr { expr, alias } => {
                    let column_name = if let Some(alias) = alias {
                        try!(new_identifier(&alias))
                    } else {
                        // if the expression is a simple identifier, make that
                        // the column name. else, assign an arbitrary name.
                        if let &ast::Expression::Ident(ref n) = &expr {
                            try!(new_identifier(n))
                        } else {
                            arbitrary_column_name()
                        }
                    };

                    let e = try!(self.ast_expression_to_sexpression(expr, &scope));
                    a.push((column_name, e));
                }
            }
        }

        Ok(a.into_iter().unzip())
    }

    fn ast_expression_to_sexpression<'b>(&mut self, ast: ast::Expression, scope: &'b SourceScope<'b>)
    -> Result<SExpression<'a, DB>, QueryPlanCompileError>
    {
        use std::borrow::IntoCow;

        match ast {
            ast::Expression::Ident(s) => {
                let column_identifier = try!(new_identifier(&s));

                let (source_id, column_offset) = match scope.get_column_offset(&column_identifier) {
                    Some(v) => v,
                    None => return Err(QueryPlanCompileError::AmbiguousColumnName(column_identifier))
                };

                Ok(SExpression::ColumnField {
                    source_id: source_id,
                    column_offset: column_offset
                })
            },
            ast::Expression::IdentMember(s1, s2) => {
                let table_identifier = try!(new_identifier(&s1));
                let column_identifier = try!(new_identifier(&s2));

                let (source_id, column_offset) = match scope.get_table_column_offset(&table_identifier, &column_identifier) {
                    Some(v) => v,
                    None => return Err(QueryPlanCompileError::AmbiguousColumnName(column_identifier))
                };

                Ok(SExpression::ColumnField {
                    source_id: source_id,
                    column_offset: column_offset
                })
            },
            ast::Expression::BinaryOp { lhs, rhs, op } => {
                let l = try!(self.ast_expression_to_sexpression(*lhs, scope));
                let r = try!(self.ast_expression_to_sexpression(*rhs, scope));

                Ok(SExpression::BinaryOp {
                    op: ast_binaryop_to_sexpression_binaryop(op),
                    lhs: Box::new(l),
                    rhs: Box::new(r)
                })
            },
            ast::Expression::StringLiteral(s) => {
                match DB::ColumnValue::from_string_literal(s.into_cow()) {
                    Ok(value) => Ok(SExpression::Value(value)),
                    Err(s) => Err(QueryPlanCompileError::BadStringLiteral(s.into_owned()))
                }
            },
            ast::Expression::Number(s) => {
                match DB::ColumnValue::from_number_literal(s.into_cow()) {
                    Ok(value) => Ok(SExpression::Value(value)),
                    Err(s) => Err(QueryPlanCompileError::BadNumberLiteral(s.into_owned()))
                }
            },
            ast::Expression::Subquery(subquery) => {
                let source_id = self.new_source_id();

                let plan = try!(self.compile_select_recurse(*subquery, scope));

                Ok(SExpression::Map {
                    source_id: source_id,
                    yield_in_fn: Box::new(plan.expr),
                    yield_out_fn: Box::new(SExpression::ColumnField {
                        source_id: source_id,
                        column_offset: 0
                    })
                })
            },
            ast::Expression::FunctionCall { name, arguments } => {
                let ident = try!(new_identifier(&name));

                macro_rules! aggregate {
                    ($op:expr) => (
                        if arguments.len() != 1 {
                            Err(QueryPlanCompileError::AggregateFunctionRequiresOneArgument)
                        } else {
                            let arg = arguments.into_iter().nth(0).unwrap();
                            let value = try!(self.ast_expression_to_sexpression(arg, scope));
                            let source_id = unimplemented!();
                            Ok(SExpression::AggregateOp {
                                op: $op,
                                source_id: source_id,
                                value: Box::new(value)
                            })
                        }
                    )
                }

                match &ident as &str {
                    "count" => aggregate!(AggregateOp::Count),
                    "avg" => aggregate!(AggregateOp::Avg),
                    _ => Err(QueryPlanCompileError::UnknownFunctionName(ident))
                }
            },
            e => panic!("unimplemented: {:?}", e)
        }
    }
}

fn ast_binaryop_to_sexpression_binaryop(ast: ast::BinaryOp) -> BinaryOp {
    match ast {
        ast::BinaryOp::Equal => BinaryOp::Equal,
        ast::BinaryOp::NotEqual => BinaryOp::NotEqual,
        ast::BinaryOp::LessThan => BinaryOp::LessThan,
        ast::BinaryOp::LessThanOrEqual => BinaryOp::LessThanOrEqual,
        ast::BinaryOp::GreaterThan => BinaryOp::GreaterThan,
        ast::BinaryOp::GreaterThanOrEqual => BinaryOp::GreaterThanOrEqual,
        ast::BinaryOp::And => BinaryOp::And,
        ast::BinaryOp::Or => BinaryOp::Or,
        ast::BinaryOp::Add => BinaryOp::Add,
        ast::BinaryOp::Subtract => BinaryOp::Subtract,
        ast::BinaryOp::Multiply => BinaryOp::Multiply,
        ast::BinaryOp::BitAnd => BinaryOp::BitAnd,
        ast::BinaryOp::BitOr => BinaryOp::BitOr,
        ast::BinaryOp::Concatenate => BinaryOp::Concatenate,
    }
}

impl<'a, DB: DatabaseInfo> fmt::Display for QueryPlan<'a, DB>
where <DB as DatabaseInfo>::Table: 'a
{
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let cn: Vec<_> = self.out_column_names.iter().map(|n| format!("`{}`", n)).collect();

        try!(writeln!(f, "query plan"));
        try!(writeln!(f, "column names: ({})", cn.connect(", ")));
        self.expr.fmt(f)
    }
}
<|endoftext|>


<commit_before>//! Matrix math is used to render things on a computer screen. This is common
//! throughout computer graphics, for examples and primers on using matrix
//! math to render things to screens please read an OpengGL tutorial.
//!
//! In wlroots we primarily use a 3x3 matrix of 32 bit floating point values to
//! represent a 2D screen. We also provide basic helper functions to assist in
//! transforming the matrices.

use Area;
use wlroots_sys::{wl_output_transform, wlr_matrix_identity, wlr_matrix_multiply,
                  wlr_matrix_project_box, wlr_matrix_projection, wlr_matrix_rotate,
                  wlr_matrix_scale, wlr_matrix_transform, wlr_matrix_translate,
                  wlr_matrix_transpose};

/// Modifies the matrix to become the identity matrix.
pub fn matrix_identity(output: &mut [f32; 9]) {
    unsafe { wlr_matrix_identity(output.as_mut_ptr()) }
}

/// Translate the matrix in x, and y.
pub fn matrix_translate(x: f32, y: f32) -> [f32; 9] {
    let mut output = [0.0; 9];
    unsafe {
        wlr_matrix_translate(output.as_mut_ptr(), x, y);
    }
    output
}

/// Scale the output in the x, and y.
pub fn matrix_scale(x: f32, y: f32) -> [f32; 9] {
    let mut output = [0.0; 9];
    unsafe {
        wlr_matrix_scale(output.as_mut_ptr(), x, y);
    }
    output
}

/// Rotate the matrix by some amount of radians.
pub fn matrix_rotate(mut matrix: [f32; 9], radians: f32) -> [f32; 9] {
    unsafe {
        wlr_matrix_rotate(matrix.as_mut_ptr(), radians);
    }
    matrix
}

/// Multiply two matrices together.
pub fn matrix_multiply(x: [f32; 9], y: [f32; 9]) -> [f32; 9] {
    let mut output = [0.0; 9];
    unsafe {
        wlr_matrix_multiply(output.as_mut_ptr(), x.as_ptr(), y.as_ptr());
    }
    output
}

/// Transform the matrix based on the given Wayland output transform mode.
pub fn matrix_transform(mut matrix: [f32; 9], transform: wl_output_transform) -> [f32; 9] {
    unsafe {
        wlr_matrix_transform(matrix.as_mut_ptr(), transform);
    }
    matrix
}

/// Transform the matrix based on the given Wayland output transform mode and
/// the width and height of a texture.
pub fn matrix_projection(mut matrix: [f32; 9],
                         width: i32,
                         height: i32,
                         transform: wl_output_transform)
                         -> [f32; 9] {
    unsafe {
        wlr_matrix_projection(matrix.as_mut_ptr(), width, height, transform);
    }
    matrix
}

pub fn matrix_transpose(matrix: [f32; 9]) -> [f32; 9] {
    let mut result = [0.0; 9];
    unsafe {
        wlr_matrix_transpose(result.as_mut_ptr(), matrix.as_ptr());
    }
    result
}

pub fn project_box(area: Area,
                   transform: wl_output_transform,
                   rotation: f32,
                   projection: [f32; 9])
                   -> [f32; 9] {
    unsafe {
        let mut output = [0.0; 9];
        wlr_matrix_project_box(output.as_mut_ptr(),
                               &area.into(),
                               transform,
                               rotation,
                               projection.as_ptr());
        output
    }
}
<commit_msg>Make matrix identity a constant<commit_after>//! Matrix math is used to render things on a computer screen. This is common
//! throughout computer graphics, for examples and primers on using matrix
//! math to render things to screens please read an OpengGL tutorial.
//!
//! In wlroots we primarily use a 3x3 matrix of 32 bit floating point values to
//! represent a 2D screen. We also provide basic helper functions to assist in
//! transforming the matrices.

use Area;
use wlroots_sys::{wl_output_transform, wlr_matrix_multiply,
                  wlr_matrix_project_box, wlr_matrix_projection, wlr_matrix_rotate,
                  wlr_matrix_scale, wlr_matrix_transform, wlr_matrix_translate,
                  wlr_matrix_transpose};

pub const IDENTITY: [f32; 9] =[
    1.0, 0.0, 0.0,
    0.0, 1.0, 0.0,
    0.0, 0.0, 1.0
];

/// Translate the matrix in x, and y.
pub fn matrix_translate(x: f32, y: f32) -> [f32; 9] {
    let mut output = [0.0; 9];
    unsafe {
        wlr_matrix_translate(output.as_mut_ptr(), x, y);
    }
    output
}

/// Scale the output in the x, and y.
pub fn matrix_scale(x: f32, y: f32) -> [f32; 9] {
    let mut output = [0.0; 9];
    unsafe {
        wlr_matrix_scale(output.as_mut_ptr(), x, y);
    }
    output
}

/// Rotate the matrix by some amount of radians.
pub fn matrix_rotate(mut matrix: [f32; 9], radians: f32) -> [f32; 9] {
    unsafe {
        wlr_matrix_rotate(matrix.as_mut_ptr(), radians);
    }
    matrix
}

/// Multiply two matrices together.
pub fn matrix_multiply(x: [f32; 9], y: [f32; 9]) -> [f32; 9] {
    let mut output = [0.0; 9];
    unsafe {
        wlr_matrix_multiply(output.as_mut_ptr(), x.as_ptr(), y.as_ptr());
    }
    output
}

/// Transform the matrix based on the given Wayland output transform mode.
pub fn matrix_transform(mut matrix: [f32; 9], transform: wl_output_transform) -> [f32; 9] {
    unsafe {
        wlr_matrix_transform(matrix.as_mut_ptr(), transform);
    }
    matrix
}

/// Transform the matrix based on the given Wayland output transform mode and
/// the width and height of a texture.
pub fn matrix_projection(mut matrix: [f32; 9],
                         width: i32,
                         height: i32,
                         transform: wl_output_transform)
                         -> [f32; 9] {
    unsafe {
        wlr_matrix_projection(matrix.as_mut_ptr(), width, height, transform);
    }
    matrix
}

pub fn matrix_transpose(matrix: [f32; 9]) -> [f32; 9] {
    let mut result = [0.0; 9];
    unsafe {
        wlr_matrix_transpose(result.as_mut_ptr(), matrix.as_ptr());
    }
    result
}

pub fn project_box(area: Area,
                   transform: wl_output_transform,
                   rotation: f32,
                   projection: [f32; 9])
                   -> [f32; 9] {
    unsafe {
        let mut output = [0.0; 9];
        wlr_matrix_project_box(output.as_mut_ptr(),
                               &area.into(),
                               transform,
                               rotation,
                               projection.as_ptr());
        output
    }
}
<|endoftext|>