<commit_before>use std::{
    net::{Ipv4Addr, SocketAddrV4, TcpListener, TcpStream},
    path::PathBuf,
};

use clap::Parser;
use dicom_core::{dicom_value, DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_encoding::transfer_syntax::TransferSyntaxIndex;
use dicom_object::{FileMetaTableBuilder, InMemDicomObject, StandardDataDictionary};
use dicom_transfer_syntax_registry::TransferSyntaxRegistry;
use dicom_ul::{pdu::PDataValueType, Pdu};
use tracing::{debug, error, info, warn, Level};

use crate::transfer::{ABSTRACT_SYNTAXES, TRANSFER_SYNTAXES};

mod transfer;

/// DICOM C-STORE SCP
#[derive(Debug, Parser)]
struct App {
    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling Application Entity title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCP")]
    calling_ae_title: String,
    /// enforce max pdu length
    #[clap(short = 's', long = "strict")]
    strict: bool,
    /// max pdu length
    #[clap(short = 'm', long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
    /// output directory for incoming objects
    #[clap(short = 'o', default_value = ".")]
    out_dir: PathBuf,
    /// Which port to listen on
    #[clap(short, default_value = "11111")]
    port: u16,
}

fn run(
    scu_stream: TcpStream,
    out_dir: PathBuf,
    calling_ae_title: &str,
    strict: bool,
    verbose: bool,
    max_pdu_length: u32,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut buffer: Vec<u8> = Vec::with_capacity(max_pdu_length as usize);
    let mut instance_buffer: Vec<u8> = Vec::with_capacity(1024 * 1024);
    let mut msgid = 1;
    let mut sop_class_uid = "".to_string();
    let mut sop_instance_uid = "".to_string();

    let mut options = dicom_ul::association::ServerAssociationOptions::new()
        .accept_any()
        .ae_title(calling_ae_title)
        .strict(strict);

    for uid in TRANSFER_SYNTAXES {
        options = options.with_transfer_syntax(*uid);
    }

    for uid in ABSTRACT_SYNTAXES {
        options = options.with_abstract_syntax(*uid);
    }

    let mut association = options.establish(scu_stream)?;

    info!("New association from {}", association.client_ae_title());
    debug!("> Presentation contexts: {:?}", association.presentation_contexts());

    loop {
        match association.receive() {
            Ok(mut pdu) => {
                if verbose {
                    debug!("scu ----> scp: {}", pdu.short_description());
                }
                match pdu {
                    Pdu::PData { ref mut data } => {
                        if data[0].value_type == PDataValueType::Data && !data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                        } else if data[0].value_type == PDataValueType::Command && data[0].is_last {
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();
                            let data_value = &data[0];
                            let v = &data_value.data;

                            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)?;
                            msgid = obj.element(tags::MESSAGE_ID)?.to_int()?;
                            sop_class_uid = obj
                                .element(tags::AFFECTED_SOP_CLASS_UID)?
                                .to_str()?
                                .to_string();
                            sop_instance_uid = obj
                                .element(tags::AFFECTED_SOP_INSTANCE_UID)?
                                .to_str()?
                                .to_string();
                            instance_buffer.clear();
                        } else if data[0].value_type == PDataValueType::Data && data[0].is_last {
                            instance_buffer.append(&mut data[0].data);

                            let presentation_context = association
                                .presentation_contexts()
                                .iter()
                                .filter(|pc| pc.id == data[0].presentation_context_id)
                                .next()
                                .unwrap();
                            let ts = &presentation_context.transfer_syntax;

                            let obj = InMemDicomObject::read_dataset_with_ts(
                                instance_buffer.as_slice(),
                                TransferSyntaxRegistry.get(ts).unwrap(),
                            )?;
                            let file_meta = FileMetaTableBuilder::new()
                                .media_storage_sop_class_uid(
                                    obj.element_by_name("SOPClassUID")?.to_str()?,
                                )
                                .media_storage_sop_instance_uid(
                                    obj.element_by_name("SOPInstanceUID")?.to_str()?,
                                )
                                .transfer_syntax(ts)
                                .build()?;
                            let file_obj = obj.with_exact_meta(file_meta);

                            // write the files to the current directory with their SOPInstanceUID as filenames
                            let mut file_path = out_dir.clone();
                            file_path
                                .push(sop_instance_uid.trim_end_matches('\0').to_string() + ".dcm");
                            file_obj.write_to_file(&file_path)?;
                            info!("Stored {}", file_path.display());

                            // send C-STORE-RSP object
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();

                            let obj =
                                create_cstore_response(msgid, &sop_class_uid, &sop_instance_uid);

                            let mut obj_data = Vec::new();

                            obj.write_dataset_with_ts(&mut obj_data, &ts)?;

                            let pdu_response = Pdu::PData {
                                data: vec![dicom_ul::pdu::PDataValue {
                                    presentation_context_id: data[0].presentation_context_id,
                                    value_type: PDataValueType::Command,
                                    is_last: true,
                                    data: obj_data,
                                }],
                            };
                            association.send(&pdu_response)?;
                        }
                    }
                    Pdu::ReleaseRQ => {
                        buffer.clear();
                        association.send(&Pdu::ReleaseRP)?;
                        info!(
                            "Released association with {}",
                            association.client_ae_title()
                        );
                    }
                    _ => {}
                }
            }
            Err(err @ dicom_ul::association::server::Error::Receive { .. }) => {
                debug!("{}", err);
                break;
            }
            Err(err) => {
                warn!("Unexpected error: {}", err);
                break;
            }
        }
    }
    info!("Association with {} dropped", association.client_ae_title());
    Ok(())
}

fn create_cstore_response(
    message_id: u16,
    sop_class_uid: &str,
    sop_instance_uid: &str,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + sop_class_uid.len() as i32
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + sop_instance_uid.len() as i32,
        ),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, sop_class_uid),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x8001]),
    ));
    // message ID being responded to
    obj.put(DataElement::new(
        tags::MESSAGE_ID_BEING_RESPONDED_TO,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));
    // status https://dicom.nema.org/dicom/2013/output/chtml/part07/chapter_C.html
    obj.put(DataElement::new(
        tags::STATUS,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));
    // SOPInstanceUID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, sop_instance_uid),
    ));

    obj
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let App {
        verbose,
        strict,
        port,
        max_pdu_length,
        out_dir,
        calling_ae_title,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .unwrap_or_else(|e| {
        eprintln!("Could not set up global logger: {}", e);
    });

    std::fs::create_dir_all(&out_dir).unwrap_or_else(|e| {
        error!("Could not create output directory: {}", e);
        std::process::exit(-2);
    });

    let listen_addr = SocketAddrV4::new(Ipv4Addr::from(0), port);
    let listener = TcpListener::bind(&listen_addr)?;
    info!("{} listening on: tcp://{}", &calling_ae_title, listen_addr);

    for stream in listener.incoming() {
        match stream {
            Ok(scu_stream) => {
                if let Err(e) = run(
                    scu_stream,
                    out_dir.clone(),
                    &calling_ae_title,
                    strict,
                    verbose,
                    max_pdu_length,
                ) {
                    error!("{}", e);
                }
            }
            Err(e) => {
                error!("{}", e);
            }
        }
    }

    Ok(())
}
<commit_msg>[storescp] Move args to `run` and add `uncompressed_only`<commit_after>use std::{
    net::{Ipv4Addr, SocketAddrV4, TcpListener, TcpStream},
    path::PathBuf,
};

use clap::Parser;
use dicom_core::{dicom_value, DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_encoding::transfer_syntax::TransferSyntaxIndex;
use dicom_object::{FileMetaTableBuilder, InMemDicomObject, StandardDataDictionary};
use dicom_transfer_syntax_registry::TransferSyntaxRegistry;
use dicom_ul::{pdu::PDataValueType, Pdu};
use tracing::{debug, error, info, warn, Level};

use crate::transfer::{ABSTRACT_SYNTAXES, NATIVE_TRANSFER_SYNTAXES, TRANSFER_SYNTAXES};

mod transfer;

/// DICOM C-STORE SCP
#[derive(Debug, Parser)]
struct App {
    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling Application Entity title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCP")]
    calling_ae_title: String,
    /// enforce max pdu length
    #[clap(short = 's', long = "strict")]
    strict: bool,
    /// Only accept native/uncompressed transfer syntaxes
    #[clap(long)]
    uncompressed_only: bool,
    /// max pdu length
    #[clap(short = 'm', long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
    /// output directory for incoming objects
    #[clap(short = 'o', default_value = ".")]
    out_dir: PathBuf,
    /// Which port to listen on
    #[clap(short, default_value = "11111")]
    port: u16,
}

fn run(scu_stream: TcpStream, args: &App) -> Result<(), Box<dyn std::error::Error>> {
    let App {
        verbose,
        calling_ae_title,
        strict,
        uncompressed_only,
        max_pdu_length,
        out_dir,
        port: _,
    } = args;
    let verbose = *verbose;

    let mut buffer: Vec<u8> = Vec::with_capacity(*max_pdu_length as usize);
    let mut instance_buffer: Vec<u8> = Vec::with_capacity(1024 * 1024);
    let mut msgid = 1;
    let mut sop_class_uid = "".to_string();
    let mut sop_instance_uid = "".to_string();

    let mut options = dicom_ul::association::ServerAssociationOptions::new()
        .accept_any()
        .ae_title(calling_ae_title)
        .strict(*strict);

    let accepted_tss = if *uncompressed_only {
        NATIVE_TRANSFER_SYNTAXES
    } else {
        TRANSFER_SYNTAXES
    };

    for uid in accepted_tss {
        options = options.with_transfer_syntax(*uid);
    }

    for uid in ABSTRACT_SYNTAXES {
        options = options.with_abstract_syntax(*uid);
    }

    let mut association = options.establish(scu_stream)?;

    info!("New association from {}", association.client_ae_title());
    debug!(
        "> Presentation contexts: {:?}",
        association.presentation_contexts()
    );

    loop {
        match association.receive() {
            Ok(mut pdu) => {
                if verbose {
                    debug!("scu ----> scp: {}", pdu.short_description());
                }
                match pdu {
                    Pdu::PData { ref mut data } => {
                        if data[0].value_type == PDataValueType::Data && !data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                        } else if data[0].value_type == PDataValueType::Command && data[0].is_last {
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();
                            let data_value = &data[0];
                            let v = &data_value.data;

                            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)?;
                            msgid = obj.element(tags::MESSAGE_ID)?.to_int()?;
                            sop_class_uid = obj
                                .element(tags::AFFECTED_SOP_CLASS_UID)?
                                .to_str()?
                                .to_string();
                            sop_instance_uid = obj
                                .element(tags::AFFECTED_SOP_INSTANCE_UID)?
                                .to_str()?
                                .to_string();
                            instance_buffer.clear();
                        } else if data[0].value_type == PDataValueType::Data && data[0].is_last {
                            instance_buffer.append(&mut data[0].data);

                            let presentation_context = association
                                .presentation_contexts()
                                .iter()
                                .filter(|pc| pc.id == data[0].presentation_context_id)
                                .next()
                                .unwrap();
                            let ts = &presentation_context.transfer_syntax;

                            let obj = InMemDicomObject::read_dataset_with_ts(
                                instance_buffer.as_slice(),
                                TransferSyntaxRegistry.get(ts).unwrap(),
                            )?;
                            let file_meta = FileMetaTableBuilder::new()
                                .media_storage_sop_class_uid(
                                    obj.element_by_name("SOPClassUID")?.to_str()?,
                                )
                                .media_storage_sop_instance_uid(
                                    obj.element_by_name("SOPInstanceUID")?.to_str()?,
                                )
                                .transfer_syntax(ts)
                                .build()?;
                            let file_obj = obj.with_exact_meta(file_meta);

                            // write the files to the current directory with their SOPInstanceUID as filenames
                            let mut file_path = out_dir.clone();
                            file_path
                                .push(sop_instance_uid.trim_end_matches('\0').to_string() + ".dcm");
                            file_obj.write_to_file(&file_path)?;
                            info!("Stored {}", file_path.display());

                            // send C-STORE-RSP object
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();

                            let obj =
                                create_cstore_response(msgid, &sop_class_uid, &sop_instance_uid);

                            let mut obj_data = Vec::new();

                            obj.write_dataset_with_ts(&mut obj_data, &ts)?;

                            let pdu_response = Pdu::PData {
                                data: vec![dicom_ul::pdu::PDataValue {
                                    presentation_context_id: data[0].presentation_context_id,
                                    value_type: PDataValueType::Command,
                                    is_last: true,
                                    data: obj_data,
                                }],
                            };
                            association.send(&pdu_response)?;
                        }
                    }
                    Pdu::ReleaseRQ => {
                        buffer.clear();
                        association.send(&Pdu::ReleaseRP)?;
                        info!(
                            "Released association with {}",
                            association.client_ae_title()
                        );
                    }
                    _ => {}
                }
            }
            Err(err @ dicom_ul::association::server::Error::Receive { .. }) => {
                debug!("{}", err);
                break;
            }
            Err(err) => {
                warn!("Unexpected error: {}", err);
                break;
            }
        }
    }
    info!("Association with {} dropped", association.client_ae_title());
    Ok(())
}

fn create_cstore_response(
    message_id: u16,
    sop_class_uid: &str,
    sop_instance_uid: &str,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + sop_class_uid.len() as i32
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + sop_instance_uid.len() as i32,
        ),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, sop_class_uid),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x8001]),
    ));
    // message ID being responded to
    obj.put(DataElement::new(
        tags::MESSAGE_ID_BEING_RESPONDED_TO,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));
    // status https://dicom.nema.org/dicom/2013/output/chtml/part07/chapter_C.html
    obj.put(DataElement::new(
        tags::STATUS,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));
    // SOPInstanceUID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, sop_instance_uid),
    ));

    obj
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if args.verbose {
                Level::DEBUG
            } else {
                Level::INFO
            })
            .finish(),
    )
    .unwrap_or_else(|e| {
        eprintln!("Could not set up global logger: {}", e);
    });

    std::fs::create_dir_all(&args.out_dir).unwrap_or_else(|e| {
        error!("Could not create output directory: {}", e);
        std::process::exit(-2);
    });

    let listen_addr = SocketAddrV4::new(Ipv4Addr::from(0), args.port);
    let listener = TcpListener::bind(&listen_addr)?;
    info!(
        "{} listening on: tcp://{}",
        &args.calling_ae_title, listen_addr
    );

    for stream in listener.incoming() {
        match stream {
            Ok(scu_stream) => {
                if let Err(e) = run(scu_stream, &args) {
                    error!("{}", e);
                }
            }
            Err(e) => {
                error!("{}", e);
            }
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>//
// sprocketnes/gfx.rs
//
// Copyright (c) 2013 Mozilla Foundation
// Author: Patrick Walton
//

use ppu;

use core::uint::range;
use sdl::sdl::{InitTimer, InitVideo};
use sdl::sdl;
use sdl::video::{DoubleBuf, HWSurface, Surface};
use sdl::video;

const SCREEN_WIDTH: uint = 320;
const SCREEN_HEIGHT: uint = 240;

pub struct Gfx {
    screen: ~Surface
}

pub impl Gfx {
    static fn new() -> Gfx {
        sdl::init([ InitVideo, InitTimer ]);
        let screen = video::set_video_mode(SCREEN_WIDTH as int,
                                           SCREEN_HEIGHT as int,
                                           24,
                                           [ HWSurface ],
                                           [ DoubleBuf ]);
        Gfx { screen: screen.unwrap() }
    }

    fn blit(&self, ppu_screen: &([u8 * 184320])) {
        let src_stride = ppu::SCREEN_WIDTH * 3;
        let dest_stride = SCREEN_WIDTH * 3;

        do self.screen.with_lock |pixels| {
            for range(0, ppu::SCREEN_HEIGHT) |y| {
                let dest_start = y * dest_stride;
                let dest = vec::mut_view(pixels, dest_start, dest_start + dest_stride);
                let src_start = y * src_stride;
                let src = vec::view(*ppu_screen, src_start, src_start + src_stride);
                vec::bytes::copy_memory(dest, src, src_stride);
            }
        }
    }
}

<commit_msg>Fix screen size and optimize blitting a bit<commit_after>//
// sprocketnes/gfx.rs
//
// Copyright (c) 2013 Mozilla Foundation
// Author: Patrick Walton
//

use ppu;

use core::uint::range;
use sdl::sdl::{InitTimer, InitVideo};
use sdl::sdl;
use sdl::video::{DoubleBuf, HWSurface, Surface};
use sdl::video;

const SCREEN_WIDTH: uint = 256;
const SCREEN_HEIGHT: uint = 240;

pub struct Gfx {
    screen: ~Surface
}

pub impl Gfx {
    static fn new() -> Gfx {
        sdl::init([ InitVideo, InitTimer ]);
        let screen = video::set_video_mode(SCREEN_WIDTH as int,
                                           SCREEN_HEIGHT as int,
                                           24,
                                           [ HWSurface ],
                                           [ DoubleBuf ]);
        Gfx { screen: screen.unwrap() }
    }

    fn blit(&self, ppu_screen: &([u8 * 184320])) {
        do self.screen.with_lock |pixels| {
            vec::bytes::copy_memory(pixels, *ppu_screen, ppu_screen.len());
        }
    }
}

<|endoftext|>


<commit_before>//! Access to various extended system registers
use bitflags::bitflags;

/// Extended feature enable mask register
#[derive(Debug)]
pub struct XCr0;

bitflags! {
    /// Configuration flags of the XCr0 register.
    ///
    /// For MPX, [`BNDREG`](XCr0Flags::BNDREG) and [`BNDCSR`](XCr0Flags::BNDCSR) must be set/unset simultaneously.
    /// For AVX-512, [`OPMASK`](XCr0Flags::OPMASK), [`ZMM_HI256`](XCr0Flags::ZMM_HI256), and [`HI16_ZMM`](XCr0Flags::HI16_ZMM) must be set/unset simultaneously.
    pub struct XCr0Flags: u64 {
        /// Enables using the x87 FPU state
        /// with `XSAVE`/`XRSTOR`.
        ///
        /// Must be set.
        const X87 = 1;
        /// Enables using MXCSR and the XMM registers
        /// with `XSAVE`/`XRSTOR`.
        ///
        /// Must be set if [`YMM`](XCr0Flags::YMM) is set.
        const SSE = 1<<1;
        /// Enables AVX instructions and using the upper halves of the YMM registers
        /// with `XSAVE`/`XRSTOR`.
        const YMM = 1<<2;
        /// Enables MPX instructions and using the BND0-BND3 bound registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const BNDREG = 1 << 3;
        /// Enables MPX instructions and using the BNDCFGU and BNDSTATUS registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const BNDCSR = 1 << 4;
        /// Enables AVX-512 instructions and using the K0-K7 mask registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const OPMASK = 1 << 5;
        /// Enables AVX-512 instructions and using the upper halves of the lower ZMM registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const ZMM_HI256 = 1 << 6;
        /// Enables AVX-512 instructions and using the upper ZMM registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const HI16_ZMM = 1 << 7;
        /// Enables using the PKRU register
        /// with `XSAVE`/`XRSTOR`.
        const MPK = 1<<9;
        /// Enables Lightweight Profiling extensions and managing LWP state
        /// with `XSAVE`/`XRSTOR` (AMD Only).
        const LWP = 1<<62;
    }
}

#[cfg(feature = "instructions")]
mod x86_64 {
    use super::*;
    impl XCr0 {
        /// Read the current set of XCR0 flags.
        #[inline]
        pub fn read() -> XCr0Flags {
            XCr0Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw XCR0 value.
        #[inline]
        pub fn read_raw() -> u64 {
            #[cfg(feature = "inline_asm")]
            unsafe {
                let (low, high): (u32, u32);
                asm!(
                    "xgetbv",
                    in("ecx") 0,
                    out("rax") low, out("rdx") high,
                    options(nomem, nostack, preserves_flags),
                );
                (high as u64) << 32 | (low as u64)
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                crate::asm::x86_64_asm_xgetbv(0)
            }
        }

        /// Write XCR0 flags.
        ///
        /// Preserves the value of reserved fields.
        /// Panics if invalid combinations of [`XCr0Flags`] are set.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to
        /// enable features that are not supported by the architecture
        #[inline]
        pub unsafe fn write(flags: XCr0Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(XCr0Flags::all().bits());
            let new_value = reserved | flags.bits();

            assert!(flags.contains(XCr0Flags::X87), "The X87 flag must be set");
            if flags.contains(XCr0Flags::YMM) {
                assert!(
                    flags.contains(XCr0Flags::SSE),
                    "AVX/YMM cannot be enabled without enabling SSE"
                );
            }
            let mpx = XCr0Flags::BNDREG | XCr0Flags::BNDCSR;
            if flags.intersects(mpx) {
                assert!(
                    flags.contains(mpx),
                    "MPX flags XCr0.BNDREG and XCr0.BNDCSR must be set and unset together"
                );
            }
            let avx512 = XCr0Flags::OPMASK | XCr0Flags::ZMM_HI256 | XCr0Flags::HI16_ZMM;
            if flags.intersects(avx512) {
                assert!(
                    flags.contains(XCr0Flags::YMM),
                    "AVX-512 cannot be enabled without enabling AVX/YMM"
                );
                assert!(
                    flags.contains(avx512),
                    "AVX-512 flags XCR0.opmask, XCR0.ZMM_Hi256, and XCR0.Hi16_ZMM must be set and unset together"
                );
            }

            Self::write_raw(new_value);
        }

        /// Write raw XCR0 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to
        /// enable features that are not supported by the architecture
        #[inline]
        pub unsafe fn write_raw(value: u64) {
            let low = value as u32;
            let high = (value >> 32) as u32;

            #[cfg(feature = "inline_asm")]
            asm!(
                "xsetbv",
                in("ecx") 0,
                in("rax") low, in("rdx") high,
                options(nomem, nostack, preserves_flags),
            );

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_xsetbv(0, low, high);
        }
    }
}
<commit_msg>XCr0: Add AVX as alias for YMM<commit_after>//! Access to various extended system registers
use bitflags::bitflags;

/// Extended feature enable mask register
#[derive(Debug)]
pub struct XCr0;

bitflags! {
    /// Configuration flags of the XCr0 register.
    ///
    /// For MPX, [`BNDREG`](XCr0Flags::BNDREG) and [`BNDCSR`](XCr0Flags::BNDCSR) must be set/unset simultaneously.
    /// For AVX-512, [`OPMASK`](XCr0Flags::OPMASK), [`ZMM_HI256`](XCr0Flags::ZMM_HI256), and [`HI16_ZMM`](XCr0Flags::HI16_ZMM) must be set/unset simultaneously.
    pub struct XCr0Flags: u64 {
        /// Enables using the x87 FPU state
        /// with `XSAVE`/`XRSTOR`.
        ///
        /// Must be set.
        const X87 = 1;
        /// Enables using MXCSR and the XMM registers
        /// with `XSAVE`/`XRSTOR`.
        ///
        /// Must be set if [`AVX`](XCr0Flags::AVX) is set.
        const SSE = 1 << 1;
        /// Enables AVX instructions and using the upper halves of the AVX registers
        /// with `XSAVE`/`XRSTOR`.
        const AVX = 1 << 2;
        /// Alias for [`AVX`](XCr0Flags::AVX)
        #[deprecated(since = "0.14.5", note = "use `AVX` instead")]
        const YMM = 1<<2;
        /// Enables MPX instructions and using the BND0-BND3 bound registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const BNDREG = 1 << 3;
        /// Enables MPX instructions and using the BNDCFGU and BNDSTATUS registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const BNDCSR = 1 << 4;
        /// Enables AVX-512 instructions and using the K0-K7 mask registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const OPMASK = 1 << 5;
        /// Enables AVX-512 instructions and using the upper halves of the lower ZMM registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const ZMM_HI256 = 1 << 6;
        /// Enables AVX-512 instructions and using the upper ZMM registers
        /// with `XSAVE`/`XRSTOR` (Intel Only).
        const HI16_ZMM = 1 << 7;
        /// Enables using the PKRU register
        /// with `XSAVE`/`XRSTOR`.
        const MPK = 1<<9;
        /// Enables Lightweight Profiling extensions and managing LWP state
        /// with `XSAVE`/`XRSTOR` (AMD Only).
        const LWP = 1<<62;
    }
}

#[cfg(feature = "instructions")]
mod x86_64 {
    use super::*;
    impl XCr0 {
        /// Read the current set of XCR0 flags.
        #[inline]
        pub fn read() -> XCr0Flags {
            XCr0Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw XCR0 value.
        #[inline]
        pub fn read_raw() -> u64 {
            #[cfg(feature = "inline_asm")]
            unsafe {
                let (low, high): (u32, u32);
                asm!(
                    "xgetbv",
                    in("ecx") 0,
                    out("rax") low, out("rdx") high,
                    options(nomem, nostack, preserves_flags),
                );
                (high as u64) << 32 | (low as u64)
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                crate::asm::x86_64_asm_xgetbv(0)
            }
        }

        /// Write XCR0 flags.
        ///
        /// Preserves the value of reserved fields.
        /// Panics if invalid combinations of [`XCr0Flags`] are set.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to
        /// enable features that are not supported by the architecture
        #[inline]
        pub unsafe fn write(flags: XCr0Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(XCr0Flags::all().bits());
            let new_value = reserved | flags.bits();

            assert!(flags.contains(XCr0Flags::X87), "The X87 flag must be set");
            if flags.contains(XCr0Flags::AVX) {
                assert!(
                    flags.contains(XCr0Flags::SSE),
                    "AVX cannot be enabled without enabling SSE"
                );
            }
            let mpx = XCr0Flags::BNDREG | XCr0Flags::BNDCSR;
            if flags.intersects(mpx) {
                assert!(
                    flags.contains(mpx),
                    "MPX flags XCr0.BNDREG and XCr0.BNDCSR must be set and unset together"
                );
            }
            let avx512 = XCr0Flags::OPMASK | XCr0Flags::ZMM_HI256 | XCr0Flags::HI16_ZMM;
            if flags.intersects(avx512) {
                assert!(
                    flags.contains(XCr0Flags::AVX),
                    "AVX-512 cannot be enabled without enabling AVX"
                );
                assert!(
                    flags.contains(avx512),
                    "AVX-512 flags XCR0.opmask, XCR0.ZMM_Hi256, and XCR0.Hi16_ZMM must be set and unset together"
                );
            }

            Self::write_raw(new_value);
        }

        /// Write raw XCR0 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to
        /// enable features that are not supported by the architecture
        #[inline]
        pub unsafe fn write_raw(value: u64) {
            let low = value as u32;
            let high = (value >> 32) as u32;

            #[cfg(feature = "inline_asm")]
            asm!(
                "xsetbv",
                in("ecx") 0,
                in("rax") low, in("rdx") high,
                options(nomem, nostack, preserves_flags),
            );

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_xsetbv(0, low, high);
        }
    }
}
<|endoftext|>


<commit_before>//! Relay for UDP implementation
//!
//! ## ShadowSocks UDP protocol
//!
//! SOCKS5 UDP Request
//! ```ignore
//! +----+------+------+----------+----------+----------+
//! |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +----+------+------+----------+----------+----------+
//! | 2  |  1   |  1   | Variable |    2     | Variable |
//! +----+------+------+----------+----------+----------+
//! ```
//!
//! SOCKS5 UDP Response
//! ```ignore
//! +----+------+------+----------+----------+----------+
//! |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +----+------+------+----------+----------+----------+
//! | 2  |  1   |  1   | Variable |    2     | Variable |
//! +----+------+------+----------+----------+----------+
//! ```
//!
//! shadowsocks UDP Request (before encrypted)
//! ```ignore
//! +------+----------+----------+----------+
//! | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +------+----------+----------+----------+
//! |  1   | Variable |    2     | Variable |
//! +------+----------+----------+----------+
//! ```
//!
//! shadowsocks UDP Response (before encrypted)
//! ```ignore
//! +------+----------+----------+----------+
//! | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +------+----------+----------+----------+
//! |  1   | Variable |    2     | Variable |
//! +------+----------+----------+----------+
//! ```
//!
//! shadowsocks UDP Request and Response (after encrypted)
//! ```ignore
//! +-------+--------------+
//! |   IV  |    PAYLOAD   |
//! +-------+--------------+
//! | Fixed |   Variable   |
//! +-------+--------------+
//! ```

use std::time::Duration;

mod association;
pub mod client;
mod crypto_io;
pub mod local;
#[cfg(feature = "local-redir")]
mod redir;
#[cfg(feature = "local-redir")]
mod redir_local;
pub mod server;
mod socks5_local;
#[cfg(feature = "local-tunnel")]
mod tunnel_local;
mod utils;

/// The maximum UDP payload size (defined in the original shadowsocks Python)
///
/// *I cannot find any references about why clowwindy used this value as the maximum
/// Socks5 UDP ASSOCIATE packet size. The only thing I can find is
/// [here](http://support.microsoft.com/kb/822061/)*
pub const MAXIMUM_UDP_PAYLOAD_SIZE: usize = 65536;

/// Default association expire time
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5 * 60);
<commit_msg>UDP associations' default timeout to be 15 seconds<commit_after>//! Relay for UDP implementation
//!
//! ## ShadowSocks UDP protocol
//!
//! SOCKS5 UDP Request
//! ```ignore
//! +----+------+------+----------+----------+----------+
//! |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +----+------+------+----------+----------+----------+
//! | 2  |  1   |  1   | Variable |    2     | Variable |
//! +----+------+------+----------+----------+----------+
//! ```
//!
//! SOCKS5 UDP Response
//! ```ignore
//! +----+------+------+----------+----------+----------+
//! |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +----+------+------+----------+----------+----------+
//! | 2  |  1   |  1   | Variable |    2     | Variable |
//! +----+------+------+----------+----------+----------+
//! ```
//!
//! shadowsocks UDP Request (before encrypted)
//! ```ignore
//! +------+----------+----------+----------+
//! | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +------+----------+----------+----------+
//! |  1   | Variable |    2     | Variable |
//! +------+----------+----------+----------+
//! ```
//!
//! shadowsocks UDP Response (before encrypted)
//! ```ignore
//! +------+----------+----------+----------+
//! | ATYP | DST.ADDR | DST.PORT |   DATA   |
//! +------+----------+----------+----------+
//! |  1   | Variable |    2     | Variable |
//! +------+----------+----------+----------+
//! ```
//!
//! shadowsocks UDP Request and Response (after encrypted)
//! ```ignore
//! +-------+--------------+
//! |   IV  |    PAYLOAD   |
//! +-------+--------------+
//! | Fixed |   Variable   |
//! +-------+--------------+
//! ```

use std::time::Duration;

mod association;
pub mod client;
mod crypto_io;
pub mod local;
#[cfg(feature = "local-redir")]
mod redir;
#[cfg(feature = "local-redir")]
mod redir_local;
pub mod server;
mod socks5_local;
#[cfg(feature = "local-tunnel")]
mod tunnel_local;
mod utils;

/// The maximum UDP payload size (defined in the original shadowsocks Python)
///
/// *I cannot find any references about why clowwindy used this value as the maximum
/// Socks5 UDP ASSOCIATE packet size. The only thing I can find is
/// [here](http://support.microsoft.com/kb/822061/)*
pub const MAXIMUM_UDP_PAYLOAD_SIZE: usize = 65536;

/// Default association expire time
///
/// FIXME: It is very hard to decide how long this value should be.
/// For some usecases, clients may open a socket for sending just one packet,
/// which will eventually cause too many *useless* associations are kept in the server.
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(15);
<|endoftext|>


<commit_before>use clap::ArgMatches;
use log;
use std;
use chrono;
use fern;

pub fn log_setup(matches: &ArgMatches) -> () {
    let mut verbosity: u64 = 0;
    let mut brevity: u64 = 0;
    if matches.is_present("verbose") {
        verbosity = matches.occurrences_of("verbose");
    }
    if matches.is_present("quiet") {
        brevity = matches.occurrences_of("quiet");
    }
    let i32_log_level: i32 = 3 - verbosity as i32 + brevity as i32;
    let log_level = match i32_log_level {
        n if n <= 0 => log::LevelFilter::Trace,
        1 => log::LevelFilter::Debug,
        2 => log::LevelFilter::Info,
        3 => log::LevelFilter::Warn,
        _ => log::LevelFilter::Error,
    };

    fern::Dispatch::new()
        .format(|out, message, record| {
            out.finish(format_args!(
                "{}[{}][{}] {}",
                chrono::Local::now().format("[%Y-%m-%d][%H:%M:%S]"),
                record.target(),
                record.level(),
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stdout())
        .apply()
        .unwrap();
}
<commit_msg>Beutify the rust code<commit_after>use chrono;
use clap::ArgMatches;
use fern;
use log;
use std;

pub fn log_setup(matches: &ArgMatches) -> () {
    let mut verbosity: u64 = 0;
    let mut brevity: u64 = 0;
    if matches.is_present("verbose") {
        verbosity = matches.occurrences_of("verbose");
    }
    if matches.is_present("quiet") {
        brevity = matches.occurrences_of("quiet");
    }
    let i32_log_level: i32 = 3 - verbosity as i32 + brevity as i32;
    let log_level = match i32_log_level {
        n if n <= 0 => log::LevelFilter::Trace,
        1 => log::LevelFilter::Debug,
        2 => log::LevelFilter::Info,
        3 => log::LevelFilter::Warn,
        _ => log::LevelFilter::Error,
    };

    fern::Dispatch::new()
        .format(|out, message, record| {
            out.finish(format_args!(
                "{}[{}][{}] {}",
                chrono::Local::now().format("[%Y-%m-%d][%H:%M:%S]"),
                record.target(),
                record.level(),
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stdout())
        .apply()
        .unwrap();
}
<|endoftext|>


<commit_before>use linked_hash_map::LinkedHashMap;

/// The type representing the inferred structure
///
/// A word of caution: Everything in this crate is "internal API", but for this type in particular,
/// since it is very central to how json_typegen works,
/// be prepared that major breaking changes may need to be made to this in the future.
#[non_exhaustive]
#[derive(Debug, PartialEq, Clone)]
pub enum Shape {
    /// `Bottom` represents the absence of any inference information
    Bottom,

    /// `Any` represents conflicting inference information that can not be
    /// represented by any single shape
    Any,

    /// `Optional(T)` represents that a value is nullable, or not always present
    Optional(Box<Shape>),

    /// Equivalent to `Optional(Bottom)`, `Null` represents optionality with no further information
    Null,

    Bool,
    StringT,
    Integer,
    Floating,
    VecT {
        elem_type: Box<Shape>,
    },
    Struct {
        fields: LinkedHashMap<String, Shape>,
    },
    Tuple(Vec<Shape>, u64),
    MapT {
        val_type: Box<Shape>,
    },
    Opaque(String),
}

pub fn fold_shapes(shapes: Vec<Shape>) -> Shape {
    shapes.into_iter().fold(Shape::Bottom, common_shape)
}

pub fn common_shape(a: Shape, b: Shape) -> Shape {
    if a == b {
        return a;
    }
    use self::Shape::*;
    match (a, b) {
        (a, Bottom) | (Bottom, a) => a,
        (Integer, Floating) | (Floating, Integer) => Floating,
        (a, Null) | (Null, a) => make_optional(a),
        (a, Optional(b)) | (Optional(b), a) => make_optional(common_shape(a, *b)),
        (Tuple(shapes1, n1), Tuple(shapes2, n2)) => {
            if shapes1.len() == shapes2.len() {
                let shapes: Vec<_> = shapes1
                    .into_iter()
                    .zip(shapes2.into_iter())
                    .map(|(a, b)| common_shape(a, b))
                    .collect();
                Tuple(shapes, n1 + n2)
            } else {
                VecT {
                    elem_type: Box::new(common_shape(fold_shapes(shapes1), fold_shapes(shapes2))),
                }
            }
        }
        (Tuple(shapes, _), VecT { elem_type: e1 }) | (VecT { elem_type: e1 }, Tuple(shapes, _)) => {
            VecT {
                elem_type: Box::new(common_shape(*e1, fold_shapes(shapes))),
            }
        }
        (VecT { elem_type: e1 }, VecT { elem_type: e2 }) => VecT {
            elem_type: Box::new(common_shape(*e1, *e2)),
        },
        (MapT { val_type: v1 }, MapT { val_type: v2 }) => MapT {
            val_type: Box::new(common_shape(*v1, *v2)),
        },
        (Struct { fields: f1 }, Struct { fields: f2 }) => Struct {
            fields: common_field_shapes(f1, f2),
        },
        (Opaque(t), _) | (_, Opaque(t)) => Opaque(t),
        _ => Any,
    }
}

fn make_optional(a: Shape) -> Shape {
    use self::Shape::*;
    match a {
        Null | Any | Bottom | Optional(_) => a,
        non_nullable => Optional(Box::new(non_nullable)),
    }
}

fn common_field_shapes(
    f1: LinkedHashMap<String, Shape>,
    mut f2: LinkedHashMap<String, Shape>,
) -> LinkedHashMap<String, Shape> {
    if f1 == f2 {
        return f1;
    }
    let mut unified = LinkedHashMap::new();
    for (key, val) in f1.into_iter() {
        match f2.remove(&key) {
            Some(val2) => unified.insert(key, common_shape(val, val2)),
            None => unified.insert(key, make_optional(val)),
        };
    }
    for (key, val) in f2.into_iter() {
        unified.insert(key, make_optional(val));
    }
    unified
}

#[test]
fn test_unify() {
    use self::Shape::*;
    assert_eq!(common_shape(Bool, Bool), Bool);
    assert_eq!(common_shape(Bool, Integer), Any);
    assert_eq!(common_shape(Integer, Floating), Floating);
    assert_eq!(common_shape(Null, Any), Any);
    assert_eq!(common_shape(Null, Bool), Optional(Box::new(Bool)));
    assert_eq!(
        common_shape(Null, Optional(Box::new(Integer))),
        Optional(Box::new(Integer))
    );
    assert_eq!(common_shape(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(common_shape(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(
        common_shape(Optional(Box::new(Integer)), Optional(Box::new(Floating))),
        Optional(Box::new(Floating))
    );
    assert_eq!(
        common_shape(Optional(Box::new(StringT)), Optional(Box::new(Integer))),
        Any
    );
}

#[test]
fn test_common_field_shapes() {
    use self::Shape::*;
    use crate::util::string_hashmap;
    {
        let f1 = string_hashmap! {
            "a" => Integer,
            "b" => Bool,
            "c" => Integer,
            "d" => StringT,
        };
        let f2 = string_hashmap! {
            "a" => Integer,
            "c" => Floating,
            "d" => Null,
            "e" => Any,
        };
        assert_eq!(
            common_field_shapes(f1, f2),
            string_hashmap! {
                "a" => Integer,
                "b" => Optional(Box::new(Bool)),
                "c" => Floating,
                "d" => Optional(Box::new(StringT)),
                "e" => Any,
            }
        );
    }
}
<commit_msg>perf: Use existing map when unifying two objects<commit_after>use linked_hash_map::LinkedHashMap;

/// The type representing the inferred structure
///
/// A word of caution: Everything in this crate is "internal API", but for this type in particular,
/// since it is very central to how json_typegen works,
/// be prepared that major breaking changes may need to be made to this in the future.
#[non_exhaustive]
#[derive(Debug, PartialEq, Clone)]
pub enum Shape {
    /// `Bottom` represents the absence of any inference information
    Bottom,

    /// `Any` represents conflicting inference information that can not be
    /// represented by any single shape
    Any,

    /// `Optional(T)` represents that a value is nullable, or not always present
    Optional(Box<Shape>),

    /// Equivalent to `Optional(Bottom)`, `Null` represents optionality with no further information
    Null,

    Bool,
    StringT,
    Integer,
    Floating,
    VecT {
        elem_type: Box<Shape>,
    },
    Struct {
        fields: LinkedHashMap<String, Shape>,
    },
    Tuple(Vec<Shape>, u64),
    MapT {
        val_type: Box<Shape>,
    },
    Opaque(String),
}

pub fn fold_shapes(shapes: Vec<Shape>) -> Shape {
    shapes.into_iter().fold(Shape::Bottom, common_shape)
}

pub fn common_shape(a: Shape, b: Shape) -> Shape {
    if a == b {
        return a;
    }
    use self::Shape::*;
    match (a, b) {
        (a, Bottom) | (Bottom, a) => a,
        (Integer, Floating) | (Floating, Integer) => Floating,
        (a, Null) | (Null, a) => make_optional(a),
        (a, Optional(b)) | (Optional(b), a) => make_optional(common_shape(a, *b)),
        (Tuple(shapes1, n1), Tuple(shapes2, n2)) => {
            if shapes1.len() == shapes2.len() {
                let shapes: Vec<_> = shapes1
                    .into_iter()
                    .zip(shapes2.into_iter())
                    .map(|(a, b)| common_shape(a, b))
                    .collect();
                Tuple(shapes, n1 + n2)
            } else {
                VecT {
                    elem_type: Box::new(common_shape(fold_shapes(shapes1), fold_shapes(shapes2))),
                }
            }
        }
        (Tuple(shapes, _), VecT { elem_type: e1 }) | (VecT { elem_type: e1 }, Tuple(shapes, _)) => {
            VecT {
                elem_type: Box::new(common_shape(*e1, fold_shapes(shapes))),
            }
        }
        (VecT { elem_type: e1 }, VecT { elem_type: e2 }) => VecT {
            elem_type: Box::new(common_shape(*e1, *e2)),
        },
        (MapT { val_type: v1 }, MapT { val_type: v2 }) => MapT {
            val_type: Box::new(common_shape(*v1, *v2)),
        },
        (Struct { fields: f1 }, Struct { fields: f2 }) => Struct {
            fields: common_field_shapes(f1, f2),
        },
        (Opaque(t), _) | (_, Opaque(t)) => Opaque(t),
        _ => Any,
    }
}

fn make_optional(a: Shape) -> Shape {
    use self::Shape::*;
    match a {
        Null | Any | Bottom | Optional(_) => a,
        non_nullable => Optional(Box::new(non_nullable)),
    }
}

fn common_field_shapes(
    mut f1: LinkedHashMap<String, Shape>,
    mut f2: LinkedHashMap<String, Shape>,
) -> LinkedHashMap<String, Shape> {
    if f1 == f2 {
        return f1;
    }
    for (key, val) in f1.iter_mut() {
        let temp = std::mem::replace(val, Shape::Bottom);
        match f2.remove(key) {
            Some(val2) => {
                *val = common_shape(temp, val2);
            }
            None => {
                *val = make_optional(temp);
            }
        };
    }
    for (key, val) in f2.into_iter() {
        f1.insert(key, make_optional(val));
    }
    f1
}

#[test]
fn test_unify() {
    use self::Shape::*;
    assert_eq!(common_shape(Bool, Bool), Bool);
    assert_eq!(common_shape(Bool, Integer), Any);
    assert_eq!(common_shape(Integer, Floating), Floating);
    assert_eq!(common_shape(Null, Any), Any);
    assert_eq!(common_shape(Null, Bool), Optional(Box::new(Bool)));
    assert_eq!(
        common_shape(Null, Optional(Box::new(Integer))),
        Optional(Box::new(Integer))
    );
    assert_eq!(common_shape(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(common_shape(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(
        common_shape(Optional(Box::new(Integer)), Optional(Box::new(Floating))),
        Optional(Box::new(Floating))
    );
    assert_eq!(
        common_shape(Optional(Box::new(StringT)), Optional(Box::new(Integer))),
        Any
    );
}

#[test]
fn test_common_field_shapes() {
    use self::Shape::*;
    use crate::util::string_hashmap;
    {
        let f1 = string_hashmap! {
            "a" => Integer,
            "b" => Bool,
            "c" => Integer,
            "d" => StringT,
        };
        let f2 = string_hashmap! {
            "a" => Integer,
            "c" => Floating,
            "d" => Null,
            "e" => Any,
        };
        assert_eq!(
            common_field_shapes(f1, f2),
            string_hashmap! {
                "a" => Integer,
                "b" => Optional(Box::new(Bool)),
                "c" => Floating,
                "d" => Optional(Box::new(StringT)),
                "e" => Any,
            }
        );
    }
}
<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/Sieve_of_Eratosthenes
#![feature(step_by)]

use std::iter::repeat;

fn int_sqrt(n: usize) -> usize {
    (n as f64).sqrt() as usize
}

// Return the prime numbers up to limit
fn simple_sieve(limit: usize) -> Vec<usize> {
    if limit < 2 {
        return vec!();
    }

    let mut primes: Vec<bool> = repeat(true).take(limit+1).collect();

    for prime in (2..int_sqrt(limit) + 1 + 1) {
        if primes[prime] {
            for multiple in (prime * prime..limit + 1).step_by(prime) {
                primes[multiple] = false
            }
        }
    }

    (2..limit + 1).filter(|&n| primes[n]).collect()
}

#[cfg(not(test))]
fn main() {
    println!("{:?}", simple_sieve(100))
}

#[test]
fn test_basic() {
    let primes = simple_sieve(30);
    assert!(primes == [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]);
}
<commit_msg>sieve_eratosthenes: remove extra +1 from testing<commit_after>// Implements http://rosettacode.org/wiki/Sieve_of_Eratosthenes
#![feature(step_by)]

use std::iter::repeat;

fn int_sqrt(n: usize) -> usize {
    (n as f64).sqrt() as usize
}

// Return the prime numbers up to limit
fn simple_sieve(limit: usize) -> Vec<usize> {
    if limit < 2 {
        return vec!();
    }

    let mut primes: Vec<bool> = repeat(true).take(limit+1).collect();

    for prime in (2..int_sqrt(limit) + 1) {
        if primes[prime] {
            for multiple in (prime * prime..limit + 1).step_by(prime) {
                primes[multiple] = false
            }
        }
    }

    (2..limit + 1).filter(|&n| primes[n]).collect()
}

#[cfg(not(test))]
fn main() {
    println!("{:?}", simple_sieve(100))
}

#[test]
fn test_basic() {
    let primes = simple_sieve(30);
    assert!(primes == [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]);
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{error, result};
use std::fmt::Debug;
use std::cmp::Ordering;

use self::rocksdb::EngineRocksdb;
use storage::{Key, Value};
use kvproto::kvrpcpb::Context;
use kvproto::errorpb::Error as ErrorHeader;

mod rocksdb;
pub mod raftkv;

// only used for rocksdb without persistent.
pub const TEMP_DIR: &'static str = "";

const SEEK_BOUND: usize = 30;

#[derive(Debug)]
pub enum Modify {
    Delete(Key),
    Put((Key, Value)),
}

pub trait Engine: Send + Sync + Debug {
    fn get(&self, ctx: &Context, key: &Key) -> Result<Option<Value>>;
    fn write(&self, ctx: &Context, batch: Vec<Modify>) -> Result<()>;
    fn snapshot<'a>(&'a self, ctx: &Context) -> Result<Box<Snapshot + 'a>>;

    fn put(&self, ctx: &Context, key: Key, value: Value) -> Result<()> {
        self.write(ctx, vec![Modify::Put((key, value))])
    }

    fn delete(&self, ctx: &Context, key: Key) -> Result<()> {
        self.write(ctx, vec![Modify::Delete(key)])
    }

    fn iter<'a>(&'a self, ctx: &Context) -> Result<Box<Cursor + 'a>>;
}

pub trait Snapshot {
    fn get(&self, key: &Key) -> Result<Option<Value>>;
    fn iter<'a>(&'a self) -> Result<Box<Cursor + 'a>>;
}

pub trait Cursor {
    fn next(&mut self) -> bool;
    fn prev(&mut self) -> bool;
    fn seek(&mut self, key: &Key) -> Result<bool>;
    fn seek_to_first(&mut self) -> bool;
    fn seek_to_last(&mut self) -> bool;
    fn valid(&self) -> bool;

    fn key(&self) -> &[u8];
    fn value(&self) -> &[u8];

    /// Seek the specified key.
    ///
    /// This method assume the current position of cursor is
    /// around `key`, otherwise you should use `seek` instead.
    fn near_seek(&mut self, key: &Key) -> Result<bool> {
        if !self.valid() {
            return self.seek(key);
        }
        if self.key() == &**key.encoded() {
            return Ok(true);
        }
        let (nav, ord): (fn(&mut _) -> bool, Ordering) = if self.key() > &**key.encoded() {
            (Cursor::prev, Ordering::Greater)
        } else {
            (Cursor::next, Ordering::Less)
        };
        let mut cnt = 0;
        while self.key().cmp(&key.encoded()) == ord && nav(self) {
            cnt += 1;
            if cnt >= SEEK_BOUND {
                return self.seek(key);
            }
        }
        // TODO: check region range.
        Ok(self.valid())
    }

    /// Get the value of specified key.
    ///
    /// This method assume the current position of cursor is
    /// around `key`, otherwise you should `seek` first.
    fn get(&mut self, key: &Key) -> Result<Option<&[u8]>> {
        if try!(self.near_seek(key)) && self.key() == &**key.encoded() {
            Ok(Some(self.value()))
        } else {
            Ok(None)
        }
    }

    fn reverse_seek(&mut self, key: &Key) -> Result<bool> {
        if !try!(self.seek(key)) && !self.seek_to_last() {
            return Ok(false);
        }

        while self.key() >= key.encoded().as_slice() && self.prev() {
        }

        Ok(self.valid())
    }

    /// Reverse seek the specified key.
    ///
    /// This method assume the current position of cursor is
    /// around `key`, otherwise you should use `reverse_seek` instead.
    fn near_reverse_seek(&mut self, key: &Key) -> Result<bool> {
        if !try!(self.near_seek(key)) && !self.seek_to_last() {
            return Ok(false);
        }

        while self.key() >= key.encoded().as_slice() && self.prev() {
        }

        Ok(self.valid())
    }
}

#[derive(Debug, Clone, Copy)]
pub enum Dsn<'a> {
    RocksDBPath(&'a str),
    RaftKv,
}

// Now we only support RocksDB.
pub fn new_engine(dsn: Dsn) -> Result<Box<Engine>> {
    match dsn {
        Dsn::RocksDBPath(path) => {
            EngineRocksdb::new(path).map(|engine| -> Box<Engine> { Box::new(engine) })
        }
        Dsn::RaftKv => unimplemented!(),
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        Request(err: ErrorHeader) {
            from()
            description("request to underhook engine failed")
            display("{:?}", err)
        }
        Other(err: Box<error::Error + Send + Sync>) {
            from()
            cause(err.as_ref())
            description(err.description())
            display("unknown error {:?}", err)
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

#[cfg(test)]
mod tests {
    use super::*;
    use tempdir::TempDir;
    use storage::make_key;
    use util::codec::bytes;
    use kvproto::kvrpcpb::Context;

    #[test]
    fn rocksdb() {
        let dir = TempDir::new("rocksdb_test").unwrap();
        let e = new_engine(Dsn::RocksDBPath(dir.path().to_str().unwrap())).unwrap();

        get_put(e.as_ref());
        batch(e.as_ref());
        seek(e.as_ref());
    }

    fn must_put(engine: &Engine, key: &[u8], value: &[u8]) {
        engine.put(&Context::new(), make_key(key), value.to_vec()).unwrap();
    }

    fn must_delete(engine: &Engine, key: &[u8]) {
        engine.delete(&Context::new(), make_key(key)).unwrap();
    }

    fn assert_has(engine: &Engine, key: &[u8], value: &[u8]) {
        assert_eq!(engine.get(&Context::new(), &make_key(key)).unwrap().unwrap(),
                   value);
    }

    fn assert_none(engine: &Engine, key: &[u8]) {
        assert_eq!(engine.get(&Context::new(), &make_key(key)).unwrap(), None);
    }

    fn assert_seek(engine: &Engine, key: &[u8], pair: (&[u8], &[u8])) {
        let mut iter = engine.iter(&Context::new()).unwrap();
        iter.seek(&make_key(key)).unwrap();
        assert_eq!((iter.key(), iter.value()),
                   (&*bytes::encode_bytes(pair.0), pair.1));
    }

    fn get_put(engine: &Engine) {
        assert_none(engine, b"x");
        must_put(engine, b"x", b"1");
        assert_has(engine, b"x", b"1");
        must_put(engine, b"x", b"2");
        assert_has(engine, b"x", b"2");
    }

    fn batch(engine: &Engine) {
        engine.write(&Context::new(),
                   vec![Modify::Put((make_key(b"x"), b"1".to_vec())),
                        Modify::Put((make_key(b"y"), b"2".to_vec()))])
            .unwrap();
        assert_has(engine, b"x", b"1");
        assert_has(engine, b"y", b"2");

        engine.write(&Context::new(),
                   vec![Modify::Delete(make_key(b"x")), Modify::Delete(make_key(b"y"))])
            .unwrap();
        assert_none(engine, b"y");
        assert_none(engine, b"y");
    }

    fn seek(engine: &Engine) {
        must_put(engine, b"x", b"1");
        assert_seek(engine, b"x", (b"x", b"1"));
        assert_seek(engine, b"a", (b"x", b"1"));
        must_put(engine, b"z", b"2");
        assert_seek(engine, b"y", (b"z", b"2"));
        assert_seek(engine, b"x\x00", (b"z", b"2"));
        let mut iter = engine.iter(&Context::new()).unwrap();
        assert!(!iter.seek(&make_key(b"z\x00")).unwrap());
        must_delete(engine, b"x");
        must_delete(engine, b"z");
    }
}
<commit_msg>storage: fix near_seek (#649)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{error, result};
use std::fmt::Debug;
use std::cmp::Ordering;

use self::rocksdb::EngineRocksdb;
use storage::{Key, Value};
use kvproto::kvrpcpb::Context;
use kvproto::errorpb::Error as ErrorHeader;

mod rocksdb;
pub mod raftkv;

// only used for rocksdb without persistent.
pub const TEMP_DIR: &'static str = "";

const SEEK_BOUND: usize = 30;

#[derive(Debug)]
pub enum Modify {
    Delete(Key),
    Put((Key, Value)),
}

pub trait Engine: Send + Sync + Debug {
    fn get(&self, ctx: &Context, key: &Key) -> Result<Option<Value>>;
    fn write(&self, ctx: &Context, batch: Vec<Modify>) -> Result<()>;
    fn snapshot<'a>(&'a self, ctx: &Context) -> Result<Box<Snapshot + 'a>>;

    fn put(&self, ctx: &Context, key: Key, value: Value) -> Result<()> {
        self.write(ctx, vec![Modify::Put((key, value))])
    }

    fn delete(&self, ctx: &Context, key: Key) -> Result<()> {
        self.write(ctx, vec![Modify::Delete(key)])
    }

    fn iter<'a>(&'a self, ctx: &Context) -> Result<Box<Cursor + 'a>>;
}

pub trait Snapshot {
    fn get(&self, key: &Key) -> Result<Option<Value>>;
    fn iter<'a>(&'a self) -> Result<Box<Cursor + 'a>>;
}

pub trait Cursor {
    fn next(&mut self) -> bool;
    fn prev(&mut self) -> bool;
    fn seek(&mut self, key: &Key) -> Result<bool>;
    fn seek_to_first(&mut self) -> bool;
    fn seek_to_last(&mut self) -> bool;
    fn valid(&self) -> bool;

    fn key(&self) -> &[u8];
    fn value(&self) -> &[u8];

    /// Seek the specified key.
    ///
    /// This method assume the current position of cursor is
    /// around `key`, otherwise you should use `seek` instead.
    fn near_seek(&mut self, key: &Key) -> Result<bool> {
        if !self.valid() {
            return self.seek(key);
        }
        if self.key() == &**key.encoded() {
            return Ok(true);
        }
        let (nav, ord): (fn(&mut _) -> bool, Ordering) = if self.key() > &**key.encoded() {
            (Cursor::prev, Ordering::Greater)
        } else {
            (Cursor::next, Ordering::Less)
        };
        let mut cnt = 0;
        while self.key().cmp(&key.encoded()) == ord && nav(self) {
            cnt += 1;
            if cnt >= SEEK_BOUND {
                return self.seek(key);
            }
        }
        if ord == Ordering::Greater {
            if self.valid() {
                if self.key() < &**key.encoded() {
                    self.next();
                }
            } else {
                self.seek_to_first();
            }
        }
        // TODO: check region range.
        Ok(self.valid())
    }

    /// Get the value of specified key.
    ///
    /// This method assume the current position of cursor is
    /// around `key`, otherwise you should `seek` first.
    fn get(&mut self, key: &Key) -> Result<Option<&[u8]>> {
        if try!(self.near_seek(key)) && self.key() == &**key.encoded() {
            Ok(Some(self.value()))
        } else {
            Ok(None)
        }
    }

    fn reverse_seek(&mut self, key: &Key) -> Result<bool> {
        if !try!(self.seek(key)) && !self.seek_to_last() {
            return Ok(false);
        }

        while self.key() >= key.encoded().as_slice() && self.prev() {
        }

        Ok(self.valid())
    }

    /// Reverse seek the specified key.
    ///
    /// This method assume the current position of cursor is
    /// around `key`, otherwise you should use `reverse_seek` instead.
    fn near_reverse_seek(&mut self, key: &Key) -> Result<bool> {
        if !try!(self.near_seek(key)) && !self.seek_to_last() {
            return Ok(false);
        }

        while self.key() >= key.encoded().as_slice() && self.prev() {
        }

        Ok(self.valid())
    }
}

#[derive(Debug, Clone, Copy)]
pub enum Dsn<'a> {
    RocksDBPath(&'a str),
    RaftKv,
}

// Now we only support RocksDB.
pub fn new_engine(dsn: Dsn) -> Result<Box<Engine>> {
    match dsn {
        Dsn::RocksDBPath(path) => {
            EngineRocksdb::new(path).map(|engine| -> Box<Engine> { Box::new(engine) })
        }
        Dsn::RaftKv => unimplemented!(),
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        Request(err: ErrorHeader) {
            from()
            description("request to underhook engine failed")
            display("{:?}", err)
        }
        Other(err: Box<error::Error + Send + Sync>) {
            from()
            cause(err.as_ref())
            description(err.description())
            display("unknown error {:?}", err)
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

#[cfg(test)]
mod tests {
    use super::*;
    use tempdir::TempDir;
    use storage::make_key;
    use util::codec::bytes;
    use util::escape;
    use kvproto::kvrpcpb::Context;

    #[test]
    fn rocksdb() {
        let dir = TempDir::new("rocksdb_test").unwrap();
        let e = new_engine(Dsn::RocksDBPath(dir.path().to_str().unwrap())).unwrap();

        get_put(e.as_ref());
        batch(e.as_ref());
        seek(e.as_ref());
        near_seek(e.as_ref());
    }

    fn must_put(engine: &Engine, key: &[u8], value: &[u8]) {
        engine.put(&Context::new(), make_key(key), value.to_vec()).unwrap();
    }

    fn must_delete(engine: &Engine, key: &[u8]) {
        engine.delete(&Context::new(), make_key(key)).unwrap();
    }

    fn assert_has(engine: &Engine, key: &[u8], value: &[u8]) {
        assert_eq!(engine.get(&Context::new(), &make_key(key)).unwrap().unwrap(),
                   value);
    }

    fn assert_none(engine: &Engine, key: &[u8]) {
        assert_eq!(engine.get(&Context::new(), &make_key(key)).unwrap(), None);
    }

    fn assert_seek(engine: &Engine, key: &[u8], pair: (&[u8], &[u8])) {
        let mut iter = engine.iter(&Context::new()).unwrap();
        iter.seek(&make_key(key)).unwrap();
        assert_eq!((iter.key(), iter.value()),
                   (&*bytes::encode_bytes(pair.0), pair.1));
    }

    fn assert_near_seek(cursor: &mut Cursor, key: &[u8], pair: (&[u8], &[u8])) {
        assert!(cursor.near_seek(&make_key(key)).unwrap(), escape(key));
        assert_eq!((cursor.key(), cursor.value()),
                   (&*bytes::encode_bytes(pair.0), pair.1));
    }

    fn assert_near_reverse_seek(cursor: &mut Cursor, key: &[u8], pair: (&[u8], &[u8])) {
        assert!(cursor.near_reverse_seek(&make_key(key)).unwrap(),
                escape(key));
        assert_eq!((cursor.key(), cursor.value()),
                   (&*bytes::encode_bytes(pair.0), pair.1));
    }

    fn get_put(engine: &Engine) {
        assert_none(engine, b"x");
        must_put(engine, b"x", b"1");
        assert_has(engine, b"x", b"1");
        must_put(engine, b"x", b"2");
        assert_has(engine, b"x", b"2");
    }

    fn batch(engine: &Engine) {
        engine.write(&Context::new(),
                   vec![Modify::Put((make_key(b"x"), b"1".to_vec())),
                        Modify::Put((make_key(b"y"), b"2".to_vec()))])
            .unwrap();
        assert_has(engine, b"x", b"1");
        assert_has(engine, b"y", b"2");

        engine.write(&Context::new(),
                   vec![Modify::Delete(make_key(b"x")), Modify::Delete(make_key(b"y"))])
            .unwrap();
        assert_none(engine, b"y");
        assert_none(engine, b"y");
    }

    fn seek(engine: &Engine) {
        must_put(engine, b"x", b"1");
        assert_seek(engine, b"x", (b"x", b"1"));
        assert_seek(engine, b"a", (b"x", b"1"));
        must_put(engine, b"z", b"2");
        assert_seek(engine, b"y", (b"z", b"2"));
        assert_seek(engine, b"x\x00", (b"z", b"2"));
        let mut iter = engine.iter(&Context::new()).unwrap();
        assert!(!iter.seek(&make_key(b"z\x00")).unwrap());
        must_delete(engine, b"x");
        must_delete(engine, b"z");
    }

    fn near_seek(engine: &Engine) {
        must_put(engine, b"x", b"1");
        must_put(engine, b"z", b"2");
        let mut cursor = engine.iter(&Context::new()).unwrap();
        let cursor_mut = cursor.as_mut();
        assert_near_seek(cursor_mut, b"x", (b"x", b"1"));
        assert_near_seek(cursor_mut, b"a", (b"x", b"1"));
        assert_near_reverse_seek(cursor_mut, b"z1", (b"z", b"2"));
        assert_near_reverse_seek(cursor_mut, b"x1", (b"x", b"1"));
        assert_near_seek(cursor_mut, b"y", (b"z", b"2"));
        assert_near_seek(cursor_mut, b"x\x00", (b"z", b"2"));
        assert!(!cursor_mut.near_seek(&make_key(b"z\x00")).unwrap());
        must_delete(engine, b"x");
        must_delete(engine, b"z");
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2019 Fabian Schuiki

//! Lowering of AST nodes to HIR nodes.

use crate::{ast_map::AstNode, crate_prelude::*, hir::HirNode};

pub(crate) fn hir_of<'gcx>(cx: &impl Context<'gcx>, node_id: NodeId) -> Result<HirNode<'gcx>> {
    let ast = cx.ast_of(node_id)?;

    #[allow(unreachable_patterns)]
    match ast {
        AstNode::Module(m) => lower_module(cx, node_id, m),
        AstNode::Port(p) => lower_port(cx, node_id, p),
        AstNode::Type(ty) => {
            let kind = match ty.data {
                ast::VoidType => hir::TypeKind::Builtin(hir::BuiltinType::Void),
                ast::BitType => hir::TypeKind::Builtin(hir::BuiltinType::Bit),
                ast::LogicType => hir::TypeKind::Builtin(hir::BuiltinType::Logic),
                ast::NamedType(name) => hir::TypeKind::Named(Spanned::new(name.name, name.span)),
                _ => return cx.unimp_msg("lowering of", ty),
            };
            let hir = hir::Type {
                id: node_id,
                span: ty.span,
                kind: kind,
            };
            Ok(HirNode::Type(cx.arena().alloc_hir(hir)))
        }
        AstNode::InstTarget(ast) => {
            let mut named_params = vec![];
            let mut pos_params = vec![];
            let mut is_pos = true;
            for param in &ast.params {
                if let Some(name) = param.name {
                    is_pos = false;
                    named_params.push((
                        param.span,
                        Spanned::new(name.name, name.span),
                        NodeId::alloc(),
                    ));
                } else {
                    if !is_pos {
                        cx.emit(
                            DiagBuilder2::warning("positional parameters must appear before named")
                                .span(param.span)
                                .add_note(format!(
                                    "assuming this refers to argument #{}",
                                    pos_params.len() + 1
                                )),
                        );
                    }
                    pos_params.push((param.span, NodeId::alloc()));
                }
            }
            let hir = hir::InstTarget {
                id: node_id,
                name: Spanned::new(ast.target.name, ast.target.span),
                span: ast.span,
                pos_params,
                named_params,
            };
            Ok(HirNode::InstTarget(cx.arena().alloc_hir(hir)))
        }
        AstNode::Inst(inst, target_id) => {
            let hir = hir::Inst {
                id: node_id,
                name: Spanned::new(inst.name.name, inst.name.span),
                span: inst.span,
                target: target_id,
                dummy: Default::default(),
            };
            Ok(HirNode::Inst(cx.arena().alloc_hir(hir)))
        }
        AstNode::TypeParam(param, decl) => {
            let hir = hir::TypeParam {
                id: node_id,
                name: Spanned::new(decl.name.name, decl.name.span),
                span: Span::union(param.span, decl.span),
                local: param.local,
                default: decl.ty.as_ref().map(|ty| {
                    let id = cx.map_ast(AstNode::Type(ty));
                    cx.set_parent(id, cx.parent_node_id(node_id).unwrap());
                    id
                }),
            };
            Ok(HirNode::TypeParam(cx.arena().alloc_hir(hir)))
        }
        _ => cx.unimp_msg("lowering of", &ast),
    }
}

fn lower_module<'gcx>(
    cx: &impl Context<'gcx>,
    node_id: NodeId,
    ast: &'gcx ast::ModDecl,
) -> Result<HirNode<'gcx>> {
    let mut next_rib = node_id;

    // Allocate parameters.
    let mut params = Vec::new();
    for param in &ast.params {
        match param.kind {
            ast::ParamKind::Type(ref decls) => {
                for decl in decls {
                    let id = cx.map_ast(AstNode::TypeParam(param, decl));
                    cx.set_parent(id, next_rib);
                    next_rib = id;
                    params.push(id);
                }
            }
            ast::ParamKind::Value(ref decls) => {
                for decl in decls {
                    let id = cx.map_ast(AstNode::ValueParam(param, decl));
                    cx.set_parent(id, next_rib);
                    next_rib = id;
                    params.push(id);
                }
            }
        }
    }

    // Allocate ports.
    let mut ports = Vec::new();
    for port in &ast.ports {
        match *port {
            ast::Port::Named { .. } => {
                let id = cx.map_ast(AstNode::Port(port));
                cx.set_parent(id, next_rib);
                next_rib = id;
                ports.push(id);
            }
            _ => return cx.unimp(port),
        }
    }

    // Allocate items.
    let mut insts = Vec::new();
    for item in &ast.items {
        match *item {
            ast::HierarchyItem::Inst(ref inst) => {
                let target_id = cx.map_ast(AstNode::InstTarget(inst));
                cx.set_parent(target_id, next_rib);
                next_rib = target_id;
                trace!(
                    "instantiation target `{}` => {:?}",
                    inst.target.name,
                    target_id
                );
                for inst in &inst.names {
                    let inst_id = cx.map_ast(AstNode::Inst(inst, target_id));
                    trace!("instantiation `{}` => {:?}", inst.name.name, inst_id);
                    cx.set_parent(inst_id, next_rib);
                    next_rib = inst_id;
                    insts.push(inst_id);
                }
            }
            // _ => return cx.unimp_msg("lowering of", item),
            _ => (),
        }
    }

    let hir = hir::Module {
        id: node_id,
        name: Spanned::new(ast.name, ast.name_span),
        span: ast.span,
        ports: cx.arena().alloc_ids(ports),
        params: cx.arena().alloc_ids(params),
        insts: cx.arena().alloc_ids(insts),
    };
    Ok(HirNode::Module(cx.arena().alloc_hir(hir)))
}

fn lower_port<'gcx>(
    cx: &impl Context<'gcx>,
    node_id: NodeId,
    ast: &'gcx ast::Port,
) -> Result<HirNode<'gcx>> {
    let hir = match *ast {
        ast::Port::Named {
            span,
            name,
            dir,
            ref ty,
            ..
        } => hir::Port {
            id: node_id,
            name: Spanned::new(name.name, name.span),
            span: span,
            dir: dir.expect("port missing direction"),
            ty: cx.map_ast(AstNode::Type(ty)),
        },
        _ => return cx.unimp(ast),
    };
    cx.set_parent(hir.ty, cx.parent_node_id(node_id).unwrap());
    Ok(HirNode::Port(cx.arena().alloc_hir(hir)))
}
<commit_msg>svlog: track param assignment type or expr<commit_after>// Copyright (c) 2016-2019 Fabian Schuiki

//! Lowering of AST nodes to HIR nodes.

use crate::{ast_map::AstNode, crate_prelude::*, hir::HirNode};

pub(crate) fn hir_of<'gcx>(cx: &impl Context<'gcx>, node_id: NodeId) -> Result<HirNode<'gcx>> {
    let ast = cx.ast_of(node_id)?;

    #[allow(unreachable_patterns)]
    match ast {
        AstNode::Module(m) => lower_module(cx, node_id, m),
        AstNode::Port(p) => lower_port(cx, node_id, p),
        AstNode::Type(ty) => {
            let kind = match ty.data {
                ast::VoidType => hir::TypeKind::Builtin(hir::BuiltinType::Void),
                ast::BitType => hir::TypeKind::Builtin(hir::BuiltinType::Bit),
                ast::LogicType => hir::TypeKind::Builtin(hir::BuiltinType::Logic),
                ast::NamedType(name) => hir::TypeKind::Named(Spanned::new(name.name, name.span)),
                _ => return cx.unimp_msg("lowering of", ty),
            };
            let hir = hir::Type {
                id: node_id,
                span: ty.span,
                kind: kind,
            };
            Ok(HirNode::Type(cx.arena().alloc_hir(hir)))
        }
        AstNode::InstTarget(ast) => {
            let mut named_params = vec![];
            let mut pos_params = vec![];
            let mut is_pos = true;
            for param in &ast.params {
                let value_id = cx.map_ast_with_parent(AstNode::TypeOrExpr(&param.expr), node_id);
                if let Some(name) = param.name {
                    is_pos = false;
                    named_params.push((param.span, Spanned::new(name.name, name.span), value_id));
                } else {
                    if !is_pos {
                        cx.emit(
                            DiagBuilder2::warning("positional parameters must appear before named")
                                .span(param.span)
                                .add_note(format!(
                                    "assuming this refers to argument #{}",
                                    pos_params.len() + 1
                                )),
                        );
                    }
                    pos_params.push((param.span, value_id));
                }
            }
            let hir = hir::InstTarget {
                id: node_id,
                name: Spanned::new(ast.target.name, ast.target.span),
                span: ast.span,
                pos_params,
                named_params,
            };
            Ok(HirNode::InstTarget(cx.arena().alloc_hir(hir)))
        }
        AstNode::Inst(inst, target_id) => {
            let hir = hir::Inst {
                id: node_id,
                name: Spanned::new(inst.name.name, inst.name.span),
                span: inst.span,
                target: target_id,
                dummy: Default::default(),
            };
            Ok(HirNode::Inst(cx.arena().alloc_hir(hir)))
        }
        AstNode::TypeParam(param, decl) => {
            let hir = hir::TypeParam {
                id: node_id,
                name: Spanned::new(decl.name.name, decl.name.span),
                span: Span::union(param.span, decl.span),
                local: param.local,
                default: decl.ty.as_ref().map(|ty| {
                    let id = cx.map_ast(AstNode::Type(ty));
                    cx.set_parent(id, cx.parent_node_id(node_id).unwrap());
                    id
                }),
            };
            Ok(HirNode::TypeParam(cx.arena().alloc_hir(hir)))
        }
        _ => cx.unimp_msg("lowering of", &ast),
    }
}

fn lower_module<'gcx>(
    cx: &impl Context<'gcx>,
    node_id: NodeId,
    ast: &'gcx ast::ModDecl,
) -> Result<HirNode<'gcx>> {
    let mut next_rib = node_id;

    // Allocate parameters.
    let mut params = Vec::new();
    for param in &ast.params {
        match param.kind {
            ast::ParamKind::Type(ref decls) => {
                for decl in decls {
                    let id = cx.map_ast(AstNode::TypeParam(param, decl));
                    cx.set_parent(id, next_rib);
                    next_rib = id;
                    params.push(id);
                }
            }
            ast::ParamKind::Value(ref decls) => {
                for decl in decls {
                    let id = cx.map_ast(AstNode::ValueParam(param, decl));
                    cx.set_parent(id, next_rib);
                    next_rib = id;
                    params.push(id);
                }
            }
        }
    }

    // Allocate ports.
    let mut ports = Vec::new();
    for port in &ast.ports {
        match *port {
            ast::Port::Named { .. } => {
                let id = cx.map_ast(AstNode::Port(port));
                cx.set_parent(id, next_rib);
                next_rib = id;
                ports.push(id);
            }
            _ => return cx.unimp(port),
        }
    }

    // Allocate items.
    let mut insts = Vec::new();
    for item in &ast.items {
        match *item {
            ast::HierarchyItem::Inst(ref inst) => {
                let target_id = cx.map_ast(AstNode::InstTarget(inst));
                cx.set_parent(target_id, next_rib);
                next_rib = target_id;
                trace!(
                    "instantiation target `{}` => {:?}",
                    inst.target.name,
                    target_id
                );
                for inst in &inst.names {
                    let inst_id = cx.map_ast(AstNode::Inst(inst, target_id));
                    trace!("instantiation `{}` => {:?}", inst.name.name, inst_id);
                    cx.set_parent(inst_id, next_rib);
                    next_rib = inst_id;
                    insts.push(inst_id);
                }
            }
            // _ => return cx.unimp_msg("lowering of", item),
            _ => (),
        }
    }

    let hir = hir::Module {
        id: node_id,
        name: Spanned::new(ast.name, ast.name_span),
        span: ast.span,
        ports: cx.arena().alloc_ids(ports),
        params: cx.arena().alloc_ids(params),
        insts: cx.arena().alloc_ids(insts),
    };
    Ok(HirNode::Module(cx.arena().alloc_hir(hir)))
}

fn lower_port<'gcx>(
    cx: &impl Context<'gcx>,
    node_id: NodeId,
    ast: &'gcx ast::Port,
) -> Result<HirNode<'gcx>> {
    let hir = match *ast {
        ast::Port::Named {
            span,
            name,
            dir,
            ref ty,
            ..
        } => hir::Port {
            id: node_id,
            name: Spanned::new(name.name, name.span),
            span: span,
            dir: dir.expect("port missing direction"),
            ty: cx.map_ast(AstNode::Type(ty)),
        },
        _ => return cx.unimp(ast),
    };
    cx.set_parent(hir.ty, cx.parent_node_id(node_id).unwrap());
    Ok(HirNode::Port(cx.arena().alloc_hir(hir)))
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All rights reserved.
//
// 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.

use std::fs;
use std::io;
use std::io::Read;
use std::path::PathBuf;

pub enum FileIterator {
    File(fs::File),
    #[cfg(test)]
    Buf(Vec<u8>, usize),
    #[cfg(all(test, feature = "benchmarks"))]
    Iter(Box<Iterator<Item = Vec<u8>> + Send>),
}

impl FileIterator {
    pub fn new(path: &PathBuf) -> io::Result<FileIterator> {
        match fs::File::open(path) {
            Ok(f) => Ok(FileIterator::File(f)),
            Err(e) => Err(e),
        }
    }
    #[cfg(test)]
    pub fn from_bytes(contents: Vec<u8>) -> FileIterator {
        FileIterator::Buf(contents, 0)
    }

    #[cfg(all(test, feature = "benchmarks"))]
    pub fn from_iter<I>(i: Box<I>) -> FileIterator
        where I: Iterator<Item = Vec<u8>> + Send + 'static
    {
        FileIterator::Iter(i)
    }
}

impl Iterator for FileIterator {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        let chunk_size = 128 * 1024;
        match self {
            &mut FileIterator::File(ref mut f) => {
                let mut buf = vec![0u8; chunk_size];
                match f.read(&mut buf[..]) {
                    Err(_) => None,
                    Ok(size) if size == 0 => None,
                    Ok(size) => Some(buf[..size].to_vec()),
                }
            }
            #[cfg(test)]
            &mut FileIterator::Buf(ref vec, ref mut pos) => {
                use std::cmp;
                if *pos >= vec.len() {
                    None
                } else {
                    let next = &vec[*pos..cmp::min(*pos + chunk_size, vec.len())];
                    *pos += chunk_size;
                    Some(next.to_owned())
                }
            }
            #[cfg(all(test, feature = "benchmarks"))]
            &mut FileIterator::Iter(ref mut inner) => inner.next(),
        }
    }
}
<commit_msg>Use BufReader for better io efficiency.<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// 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.

use std::fs;
use std::io;
use std::io::Read;
use std::path::PathBuf;

pub enum FileIterator {
    File(io::BufReader<fs::File>),
    #[cfg(test)]
    Buf(Vec<u8>, usize),
    #[cfg(all(test, feature = "benchmarks"))]
    Iter(Box<Iterator<Item = Vec<u8>> + Send>),
}

impl FileIterator {
    pub fn new(path: &PathBuf) -> io::Result<FileIterator> {
        match fs::File::open(path) {
            Ok(f) => Ok(FileIterator::File(io::BufReader::new(f))),
            Err(e) => Err(e),
        }
    }
    #[cfg(test)]
    pub fn from_bytes(contents: Vec<u8>) -> FileIterator {
        FileIterator::Buf(contents, 0)
    }

    #[cfg(all(test, feature = "benchmarks"))]
    pub fn from_iter<I>(i: Box<I>) -> FileIterator
        where I: Iterator<Item = Vec<u8>> + Send + 'static
    {
        FileIterator::Iter(i)
    }
}

impl Iterator for FileIterator {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        let chunk_size = 128 * 1024;
        match self {
            &mut FileIterator::File(ref mut f) => {
                let mut buf = vec![0u8; chunk_size];
                match f.read(&mut buf[..]) {
                    Err(_) => None,
                    Ok(size) if size == 0 => None,
                    Ok(size) => Some(buf[..size].to_vec()),
                }
            }
            #[cfg(test)]
            &mut FileIterator::Buf(ref vec, ref mut pos) => {
                use std::cmp;
                if *pos >= vec.len() {
                    None
                } else {
                    let next = &vec[*pos..cmp::min(*pos + chunk_size, vec.len())];
                    *pos += chunk_size;
                    Some(next.to_owned())
                }
            }
            #[cfg(all(test, feature = "benchmarks"))]
            &mut FileIterator::Iter(ref mut inner) => inner.next(),
        }
    }
}
<|endoftext|>