<commit_before>use anyhow::{anyhow, Context, Result};
use async_std::sync::RwLock;
use async_std::task;
use async_std::{stream, task::JoinHandle};
use async_trait::async_trait;
use futures::StreamExt;
use homectl_types::device::CorrelatedColorTemperature;
use homectl_types::{
    device::{Device, DeviceId, DeviceState, Light},
    event::{Message, TxEventChannel},
    integration::{Integration, IntegrationActionPayload, IntegrationId},
};
use palette::Hsv;
use rust_async_tuyapi::{tuyadevice::TuyaDevice, Payload, PayloadStruct};
use serde::Deserialize;
use serde_json::{json, Value};
use std::ops::Range;
use std::{
    collections::HashMap,
    net::IpAddr,
    str::FromStr,
    sync::Arc,
    time::{Duration, SystemTime},
};

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaDeviceConfig {
    name: String,
    local_key: String,
    ip: String,
    power_on_field: String,
    brightness_field: String,
    color_temp_field: String,
}

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaConfig {
    devices: HashMap<DeviceId, TuyaDeviceConfig>,
}

pub struct Tuya {
    id: IntegrationId,
    event_tx: TxEventChannel,
    config: TuyaConfig,
    device_expected_states: HashMap<DeviceId, Arc<RwLock<Device>>>,
    device_poll_handles: HashMap<DeviceId, JoinHandle<()>>,
}

fn default_device(device_id: DeviceId, name: String, integration_id: IntegrationId) -> Device {
    Device {
        id: device_id,
        name,
        integration_id,
        scene: None,
        state: DeviceState::Light(Light {
            power: false,
            brightness: None,
            color: None,
            cct: None,
            transition_ms: None,
        }),
    }
}

#[async_trait]
impl Integration for Tuya {
    fn new(id: &IntegrationId, config: &config::Value, event_tx: TxEventChannel) -> Result<Self> {
        let config: TuyaConfig = config
            .clone()
            .try_into()
            .context("Failed to deserialize config of Tuya integration")?;

        Ok(Tuya {
            id: id.clone(),
            config,
            event_tx,
            device_expected_states: HashMap::new(),
            device_poll_handles: HashMap::new(),
        })
    }

    async fn register(&mut self) -> Result<()> {
        let mut device_expected_states = HashMap::new();

        let integration_id = self.id.clone();

        for (device_id, device_config) in &self.config.devices {
            let device_id = device_id.clone();
            let device_config = device_config.clone();
            let event_tx = self.event_tx.clone();
            let integration_id = integration_id.clone();

            println!("Getting initial state of {}", device_config.name);
            let device = get_tuya_state(&device_id, &integration_id, &device_config).await;
            let device = device.unwrap_or_else(|_| {
                println!("Failed to get initial state of Tuya device {}, creating Device with default state", device_config.name);

                default_device(device_id.clone(), device_config.name.clone(), integration_id)
            });

            device_expected_states.insert(device_id.clone(), Arc::new(RwLock::new(device.clone())));

            event_tx.send(Message::IntegrationDeviceRefresh { device });
        }

        self.device_expected_states = device_expected_states;

        println!("registered tuya integration {}", self.id);

        Ok(())
    }

    async fn start(&mut self) -> Result<()> {
        println!("starting tuya integration {}", self.id);

        let device_expected_states = self.device_expected_states.clone();

        for (device_id, device_config) in &self.config.devices {
            let device_config = device_config.clone();
            let device_expected_state = device_expected_states.get(device_id).unwrap().clone();
            let sender = self.event_tx.clone();

            let handle = task::spawn(async move {
                poll_light(&device_config, sender, device_expected_state).await
            });

            self.device_poll_handles.insert(device_id.clone(), handle);
        }

        println!("started tuya integration {}", self.id);

        Ok(())
    }

    async fn set_integration_device_state(&mut self, device: &Device) -> Result<()> {
        {
            let device_expected_state = self.device_expected_states.get(&device.id).unwrap();
            let mut device_expected_state = device_expected_state.write().await;
            *device_expected_state = device.clone();
            dbg!(device_expected_state);
        }

        let mut device_poll_handle = self.device_poll_handles.remove(&device.id);
        if let Some(poll_handle) = device_poll_handle.take() {
            poll_handle.cancel().await;
        }

        let config = self.config.clone();
        let sender = self.event_tx.clone();

        let device_config = config
            .devices
            .get(&device.id)
            .context(format!(
                "Could not find TuyaDeviceConfig for device with id {}",
                device.id,
            ))?
            .clone();
        let device_expected_state = self.device_expected_states.get(&device.id).unwrap().clone();

        let handle = {
            let device_config = device_config.clone();
            task::spawn(
                async move { poll_light(&device_config, sender, device_expected_state).await },
            )
        };

        self.device_poll_handles.insert(device.id.clone(), handle);

        set_tuya_state(device, &device_config).await?;
        Ok(())
    }

    async fn run_integration_action(&mut self, _: &IntegrationActionPayload) -> Result<()> {
        Ok(())
    }
}

#[derive(Debug)]
struct TuyaState {
    power_on: bool,
    brightness: u32,
    color_temperature: u32,
}

fn to_tuya_state(device: &Device) -> Result<TuyaState> {
    let light_state = match device.state {
        DeviceState::Light(Light {
            brightness,
            color,
            power,
            transition_ms,
            ref cct,
        }) => Ok(Light {
            power,
            brightness,
            color,
            transition_ms,
            cct: cct.clone(),
        }),
        _ => Err(anyhow!("Unsupported device state")),
    }?;

    let color = light_state.color.unwrap_or_else(|| Hsv::new(0.0, 1.0, 1.0));

    let brightness = light_state.brightness.unwrap_or(1.0) * color.value;

    // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
    let brightness = f32::floor(brightness * 990.0) as u32 + 10;

    let cct = light_state.cct.map(|cct| cct.get_cct()).unwrap_or(4000.0);

    // Range of my bulbs is from 2700K - 4100K (and they express this as a
    // 0-1000 range), this is very likely not true for all Tuya bulbs
    let min_supported_temp = 2700.0;
    let max_supported_temp = 4100.0;

    // Scale the value into 0.0 - 1.0 range
    let q = (cct - min_supported_temp) / (max_supported_temp - min_supported_temp);
    let q = q.clamp(0.0, 1.0);

    // Scale the value into 0 - 1000 range
    let color_temperature = f32::floor(q * 1000.0) as u32;

    let state = TuyaState {
        power_on: light_state.power,
        brightness,
        color_temperature,
    };

    Ok(state)
}

async fn set_tuya_state(device: &Device, device_config: &TuyaDeviceConfig) -> Result<()> {
    // println!("setting tuya state: {:?} {}", device.state, device.name);
    let tuya_state = to_tuya_state(device)?;

    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        let mut dps = HashMap::new();
        dps.insert(
            device_config.power_on_field.clone(),
            json!(tuya_state.power_on),
        );
        dps.insert(
            device_config.brightness_field.clone(),
            json!(tuya_state.brightness),
        );
        dps.insert(
            device_config.color_temp_field.clone(),
            json!(tuya_state.color_temperature),
        );

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device.id.to_string(),
            // gw_id: None,
            gw_id: Some(device.id.to_string()),
            uid: Some(device.id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(dps),
        });

        tokio::time::timeout(Duration::from_millis(250), tuya_device.set(payload, 0)).await??
    }

    Ok(())
}

async fn get_tuya_state(
    device_id: &DeviceId,
    integration_id: &IntegrationId,
    device_config: &TuyaDeviceConfig,
) -> Result<Device> {
    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    let response = {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device_id.to_string(),
            gw_id: Some(device_id.to_string()),
            uid: Some(device_id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(HashMap::new()),
        });

        tokio::time::timeout(Duration::from_millis(250), tuya_device.get(payload, 0))
            .await?
            .context(format!("Error while polling {}", device_config.name))?
    };

    let first = response
        .first()
        .context("Could not find valid Tuya Message in response")?;

    if let Payload::Struct(PayloadStruct { dps: Some(dps), .. }) = &first.payload {
        let power = if let Some(Value::Bool(value)) = dps.get(&device_config.power_on_field) {
            *value
        } else {
            true
        };

        let brightness =
            if let Some(Value::Number(value)) = dps.get(&device_config.brightness_field) {
                // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
                let bri = value.as_u64().unwrap_or(1000);

                // Scale range to 0-990
                let bri = bri - 10;

                // Scale range to 0-1
                Some(bri as f32 / 990.0)
            } else {
                None
            };

        let cct = if let Some(Value::Number(value)) = dps.get(&device_config.color_temp_field) {
            // Range of my bulbs is from 2700K - 4100K (and they express this as a
            // 0-1000 range), this is very likely not true for all Tuya bulbs
            let min_supported_temp = 2700.0;
            let max_supported_temp = 4100.0;

            let ct = value.as_u64().unwrap_or(1000);

            // Scale range to 0-1
            let q = ct as f32 / 1000.0;

            let cct = q * (max_supported_temp - min_supported_temp) + min_supported_temp;

            Some(CorrelatedColorTemperature::new(
                cct,
                Range {
                    start: min_supported_temp,
                    end: max_supported_temp,
                },
            ))
        } else {
            None
        };

        let state = DeviceState::Light(Light::new_with_cct(power, brightness, cct, Some(1000)));

        let device = Device {
            id: device_id.clone(),
            name: device_config.name.clone(),
            integration_id: integration_id.clone(),
            scene: None,
            state,
        };

        Ok(device)
    } else {
        Err(anyhow!("Unsupported Tuya response"))
    }
}

pub async fn poll_light(
    device_config: &TuyaDeviceConfig,
    sender: TxEventChannel,
    device_expected_state: Arc<RwLock<Device>>,
) {
    let poll_rate = Duration::from_millis(1000);
    let mut interval = stream::interval(poll_rate);

    loop {
        interval.next().await;

        let device_expected_state = { device_expected_state.read().await.clone() };
        let result = set_tuya_state(&device_expected_state, device_config).await;

        if let Err(e) = result {
            eprintln!(
                "Error while polling Tuya state for device {}: {:?}",
                device_expected_state.name, e
            );
        }

        // We still need to send our version of the device state to homectl core, in
        // case it has gone stale.
        sender.send(Message::IntegrationDeviceRefresh {
            device: device_expected_state,
        });
    }
}
<commit_msg>Remove debug print<commit_after>use anyhow::{anyhow, Context, Result};
use async_std::sync::RwLock;
use async_std::task;
use async_std::{stream, task::JoinHandle};
use async_trait::async_trait;
use futures::StreamExt;
use homectl_types::device::CorrelatedColorTemperature;
use homectl_types::{
    device::{Device, DeviceId, DeviceState, Light},
    event::{Message, TxEventChannel},
    integration::{Integration, IntegrationActionPayload, IntegrationId},
};
use palette::Hsv;
use rust_async_tuyapi::{tuyadevice::TuyaDevice, Payload, PayloadStruct};
use serde::Deserialize;
use serde_json::{json, Value};
use std::ops::Range;
use std::{
    collections::HashMap,
    net::IpAddr,
    str::FromStr,
    sync::Arc,
    time::{Duration, SystemTime},
};

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaDeviceConfig {
    name: String,
    local_key: String,
    ip: String,
    power_on_field: String,
    brightness_field: String,
    color_temp_field: String,
}

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaConfig {
    devices: HashMap<DeviceId, TuyaDeviceConfig>,
}

pub struct Tuya {
    id: IntegrationId,
    event_tx: TxEventChannel,
    config: TuyaConfig,
    device_expected_states: HashMap<DeviceId, Arc<RwLock<Device>>>,
    device_poll_handles: HashMap<DeviceId, JoinHandle<()>>,
}

fn default_device(device_id: DeviceId, name: String, integration_id: IntegrationId) -> Device {
    Device {
        id: device_id,
        name,
        integration_id,
        scene: None,
        state: DeviceState::Light(Light {
            power: false,
            brightness: None,
            color: None,
            cct: None,
            transition_ms: None,
        }),
    }
}

#[async_trait]
impl Integration for Tuya {
    fn new(id: &IntegrationId, config: &config::Value, event_tx: TxEventChannel) -> Result<Self> {
        let config: TuyaConfig = config
            .clone()
            .try_into()
            .context("Failed to deserialize config of Tuya integration")?;

        Ok(Tuya {
            id: id.clone(),
            config,
            event_tx,
            device_expected_states: HashMap::new(),
            device_poll_handles: HashMap::new(),
        })
    }

    async fn register(&mut self) -> Result<()> {
        let mut device_expected_states = HashMap::new();

        let integration_id = self.id.clone();

        for (device_id, device_config) in &self.config.devices {
            let device_id = device_id.clone();
            let device_config = device_config.clone();
            let event_tx = self.event_tx.clone();
            let integration_id = integration_id.clone();

            println!("Getting initial state of {}", device_config.name);
            let device = get_tuya_state(&device_id, &integration_id, &device_config).await;
            let device = device.unwrap_or_else(|_| {
                println!("Failed to get initial state of Tuya device {}, creating Device with default state", device_config.name);

                default_device(device_id.clone(), device_config.name.clone(), integration_id)
            });

            device_expected_states.insert(device_id.clone(), Arc::new(RwLock::new(device.clone())));

            event_tx.send(Message::IntegrationDeviceRefresh { device });
        }

        self.device_expected_states = device_expected_states;

        println!("registered tuya integration {}", self.id);

        Ok(())
    }

    async fn start(&mut self) -> Result<()> {
        println!("starting tuya integration {}", self.id);

        let device_expected_states = self.device_expected_states.clone();

        for (device_id, device_config) in &self.config.devices {
            let device_config = device_config.clone();
            let device_expected_state = device_expected_states.get(device_id).unwrap().clone();
            let sender = self.event_tx.clone();

            let handle = task::spawn(async move {
                poll_light(&device_config, sender, device_expected_state).await
            });

            self.device_poll_handles.insert(device_id.clone(), handle);
        }

        println!("started tuya integration {}", self.id);

        Ok(())
    }

    async fn set_integration_device_state(&mut self, device: &Device) -> Result<()> {
        {
            let device_expected_state = self.device_expected_states.get(&device.id).unwrap();
            let mut device_expected_state = device_expected_state.write().await;
            *device_expected_state = device.clone();
        }

        let mut device_poll_handle = self.device_poll_handles.remove(&device.id);
        if let Some(poll_handle) = device_poll_handle.take() {
            poll_handle.cancel().await;
        }

        let config = self.config.clone();
        let sender = self.event_tx.clone();

        let device_config = config
            .devices
            .get(&device.id)
            .context(format!(
                "Could not find TuyaDeviceConfig for device with id {}",
                device.id,
            ))?
            .clone();
        let device_expected_state = self.device_expected_states.get(&device.id).unwrap().clone();

        let handle = {
            let device_config = device_config.clone();
            task::spawn(
                async move { poll_light(&device_config, sender, device_expected_state).await },
            )
        };

        self.device_poll_handles.insert(device.id.clone(), handle);

        set_tuya_state(device, &device_config).await?;
        Ok(())
    }

    async fn run_integration_action(&mut self, _: &IntegrationActionPayload) -> Result<()> {
        Ok(())
    }
}

#[derive(Debug)]
struct TuyaState {
    power_on: bool,
    brightness: u32,
    color_temperature: u32,
}

fn to_tuya_state(device: &Device) -> Result<TuyaState> {
    let light_state = match device.state {
        DeviceState::Light(Light {
            brightness,
            color,
            power,
            transition_ms,
            ref cct,
        }) => Ok(Light {
            power,
            brightness,
            color,
            transition_ms,
            cct: cct.clone(),
        }),
        _ => Err(anyhow!("Unsupported device state")),
    }?;

    let color = light_state.color.unwrap_or_else(|| Hsv::new(0.0, 1.0, 1.0));

    let brightness = light_state.brightness.unwrap_or(1.0) * color.value;

    // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
    let brightness = f32::floor(brightness * 990.0) as u32 + 10;

    let cct = light_state.cct.map(|cct| cct.get_cct()).unwrap_or(4000.0);

    // Range of my bulbs is from 2700K - 4100K (and they express this as a
    // 0-1000 range), this is very likely not true for all Tuya bulbs
    let min_supported_temp = 2700.0;
    let max_supported_temp = 4100.0;

    // Scale the value into 0.0 - 1.0 range
    let q = (cct - min_supported_temp) / (max_supported_temp - min_supported_temp);
    let q = q.clamp(0.0, 1.0);

    // Scale the value into 0 - 1000 range
    let color_temperature = f32::floor(q * 1000.0) as u32;

    let state = TuyaState {
        power_on: light_state.power,
        brightness,
        color_temperature,
    };

    Ok(state)
}

async fn set_tuya_state(device: &Device, device_config: &TuyaDeviceConfig) -> Result<()> {
    // println!("setting tuya state: {:?} {}", device.state, device.name);
    let tuya_state = to_tuya_state(device)?;

    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        let mut dps = HashMap::new();
        dps.insert(
            device_config.power_on_field.clone(),
            json!(tuya_state.power_on),
        );
        dps.insert(
            device_config.brightness_field.clone(),
            json!(tuya_state.brightness),
        );
        dps.insert(
            device_config.color_temp_field.clone(),
            json!(tuya_state.color_temperature),
        );

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device.id.to_string(),
            // gw_id: None,
            gw_id: Some(device.id.to_string()),
            uid: Some(device.id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(dps),
        });

        tokio::time::timeout(Duration::from_millis(250), tuya_device.set(payload, 0)).await??
    }

    Ok(())
}

async fn get_tuya_state(
    device_id: &DeviceId,
    integration_id: &IntegrationId,
    device_config: &TuyaDeviceConfig,
) -> Result<Device> {
    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    let response = {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device_id.to_string(),
            gw_id: Some(device_id.to_string()),
            uid: Some(device_id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(HashMap::new()),
        });

        tokio::time::timeout(Duration::from_millis(250), tuya_device.get(payload, 0))
            .await?
            .context(format!("Error while polling {}", device_config.name))?
    };

    let first = response
        .first()
        .context("Could not find valid Tuya Message in response")?;

    if let Payload::Struct(PayloadStruct { dps: Some(dps), .. }) = &first.payload {
        let power = if let Some(Value::Bool(value)) = dps.get(&device_config.power_on_field) {
            *value
        } else {
            true
        };

        let brightness =
            if let Some(Value::Number(value)) = dps.get(&device_config.brightness_field) {
                // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
                let bri = value.as_u64().unwrap_or(1000);

                // Scale range to 0-990
                let bri = bri - 10;

                // Scale range to 0-1
                Some(bri as f32 / 990.0)
            } else {
                None
            };

        let cct = if let Some(Value::Number(value)) = dps.get(&device_config.color_temp_field) {
            // Range of my bulbs is from 2700K - 4100K (and they express this as a
            // 0-1000 range), this is very likely not true for all Tuya bulbs
            let min_supported_temp = 2700.0;
            let max_supported_temp = 4100.0;

            let ct = value.as_u64().unwrap_or(1000);

            // Scale range to 0-1
            let q = ct as f32 / 1000.0;

            let cct = q * (max_supported_temp - min_supported_temp) + min_supported_temp;

            Some(CorrelatedColorTemperature::new(
                cct,
                Range {
                    start: min_supported_temp,
                    end: max_supported_temp,
                },
            ))
        } else {
            None
        };

        let state = DeviceState::Light(Light::new_with_cct(power, brightness, cct, Some(1000)));

        let device = Device {
            id: device_id.clone(),
            name: device_config.name.clone(),
            integration_id: integration_id.clone(),
            scene: None,
            state,
        };

        Ok(device)
    } else {
        Err(anyhow!("Unsupported Tuya response"))
    }
}

pub async fn poll_light(
    device_config: &TuyaDeviceConfig,
    sender: TxEventChannel,
    device_expected_state: Arc<RwLock<Device>>,
) {
    let poll_rate = Duration::from_millis(1000);
    let mut interval = stream::interval(poll_rate);

    loop {
        interval.next().await;

        let device_expected_state = { device_expected_state.read().await.clone() };
        let result = set_tuya_state(&device_expected_state, device_config).await;

        if let Err(e) = result {
            eprintln!(
                "Error while polling Tuya state for device {}: {:?}",
                device_expected_state.name, e
            );
        }

        // We still need to send our version of the device state to homectl core, in
        // case it has gone stale.
        sender.send(Message::IntegrationDeviceRefresh {
            device: device_expected_state,
        });
    }
}
<|endoftext|>


<commit_before>//! Loggers and logging events for differential dataflow.

/// Logger for differential dataflow events.
pub type Logger = ::timely::logging::Logger<DifferentialEvent>;

/// Possible different differential events.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub enum DifferentialEvent {
    /// Batch creation.
    Batch(BatchEvent),
    /// Merge start and stop events.
    Merge(MergeEvent),
    /// A merge failed to complete in time.
    MergeShortfall(MergeShortfall),
}

/// Either the start or end of a merge event.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub struct BatchEvent {
    /// Operator identifier.
    pub operator: usize,
    /// Which order of magnitude.
    pub length: usize,
}

impl From<BatchEvent> for DifferentialEvent { fn from(e: BatchEvent) -> Self { DifferentialEvent::Batch(e) } }

/// Either the start or end of a merge event.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub struct MergeEvent {
    /// Operator identifier.
    pub operator: usize,
    /// Which order of magnitude.
    pub scale: usize,
    /// Length of first trace.
    pub length1: usize,
    /// Length of second trace.
    pub length2: usize,
    /// None implies a start.
    pub complete: Option<usize>,
}

impl From<MergeEvent> for DifferentialEvent { fn from(e: MergeEvent) -> Self { DifferentialEvent::Merge(e) } }

/// A merge failed to complete in time.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub struct MergeShortfall {
    /// Operator identifer.
    pub operator: usize,
    /// Which order of magnitude.
    pub scale: usize,
    /// By how much were we short.
    pub shortfall: usize,
}

impl From<MergeShortfall> for DifferentialEvent { fn from(e: MergeShortfall) -> Self { DifferentialEvent::MergeShortfall(e) } }
<commit_msg>enable logging function<commit_after>//! Loggers and logging events for differential dataflow.

/// Logger for differential dataflow events.
pub type Logger = ::timely::logging::Logger<DifferentialEvent>;

/// Enables logging of differential dataflow events.
pub fn enable<A, W>(worker: &mut timely::worker::Worker<A>, writer: W) -> Option<Box<std::any::Any+'static>>
where
    A: timely::communication::Allocate,
    W: std::io::Write+'static,
{
    let writer = ::timely::dataflow::operators::capture::EventWriter::new(writer);
    let mut logger = ::timely::logging::BatchLogger::new(writer);
    worker
        .log_register()
        .insert::<DifferentialEvent,_>("differential/arrange", move |time, data| logger.publish_batch(time, data))
}

/// Possible different differential events.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub enum DifferentialEvent {
    /// Batch creation.
    Batch(BatchEvent),
    /// Merge start and stop events.
    Merge(MergeEvent),
    /// A merge failed to complete in time.
    MergeShortfall(MergeShortfall),
}

/// Either the start or end of a merge event.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub struct BatchEvent {
    /// Operator identifier.
    pub operator: usize,
    /// Which order of magnitude.
    pub length: usize,
}

impl From<BatchEvent> for DifferentialEvent { fn from(e: BatchEvent) -> Self { DifferentialEvent::Batch(e) } }

/// Either the start or end of a merge event.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub struct MergeEvent {
    /// Operator identifier.
    pub operator: usize,
    /// Which order of magnitude.
    pub scale: usize,
    /// Length of first trace.
    pub length1: usize,
    /// Length of second trace.
    pub length2: usize,
    /// None implies a start.
    pub complete: Option<usize>,
}

impl From<MergeEvent> for DifferentialEvent { fn from(e: MergeEvent) -> Self { DifferentialEvent::Merge(e) } }

/// A merge failed to complete in time.
#[derive(Debug, Clone, Abomonation, Ord, PartialOrd, Eq, PartialEq)]
pub struct MergeShortfall {
    /// Operator identifer.
    pub operator: usize,
    /// Which order of magnitude.
    pub scale: usize,
    /// By how much were we short.
    pub shortfall: usize,
}

impl From<MergeShortfall> for DifferentialEvent { fn from(e: MergeShortfall) -> Self { DifferentialEvent::MergeShortfall(e) } }
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

#[macro_use]
extern crate log;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimir;
extern crate rs_es;
extern crate chrono;

extern crate mimirsbrunn;
extern crate serde;
extern crate geo;

use std::collections::{HashSet, HashMap};
use mimir::rubber::Rubber;

pub type AdminsVec = Vec<Rc<mimir::Admin>>;
pub type StreetsVec = Vec<mimir::Street>;
pub type OsmPbfReader = osmpbfreader::OsmPbfReader<std::fs::File>;

use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::rc::Rc;
use std::cell::Cell;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
    flag_import_way: bool,
    flag_import_admin: bool,
    flag_dataset: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] [--import-way] [--import-admin] [--dataset=<dataset>] --level=<level> ...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -w, --import-way             Import ways
    -a, --import-admin           Import admins
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
    -d, --dataset=<dataset>         Name of the dataset, [default: fr]
";

#[derive(Debug)]
struct AdminMatcher {
    admin_levels: HashSet<u32>,
}
impl AdminMatcher {
    pub fn new(levels: HashSet<u32>) -> AdminMatcher {
        AdminMatcher { admin_levels: levels }
    }

    pub fn is_admin(&self, obj: &osmpbfreader::OsmObj) -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("boundary").map_or(false, |v| v == "administrative") &&
                rel.tags.get("admin_level").map_or(false, |lvl| {
                    self.admin_levels.contains(&lvl.parse::<u32>().unwrap_or(0))
                })
            }
            _ => false,
        }
    }
}

fn parse_osm_pbf(path: &str) -> OsmPbfReader {
    let path = std::path::Path::new(&path);
    osmpbfreader::OsmPbfReader::new(std::fs::File::open(&path).unwrap())
}

fn administrative_regions(pbf: &mut OsmPbfReader, levels: HashSet<u32>) -> AdminsVec {
    let mut administrative_regions = AdminsVec::new();
    let matcher = AdminMatcher::new(levels);
    let objects = osmpbfreader::get_objs_and_deps(pbf, |o| matcher.is_admin(o)).unwrap();
    // load administratives regions
    for (_, obj) in &objects {
        if !matcher.is_admin(&obj) {
            continue;
        }
        if let &osmpbfreader::OsmObj::Relation(ref relation) = obj {
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    warn!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            // admininstrative region without coordinates
            let coord_centre = relation.refs
                                       .iter()
                                       .find(|rf| rf.role == "admin_centre")
                                       .and_then(|r| {
                                           objects.get(&r.member).and_then(|value| {
                                               match value {
                                                   &osmpbfreader::OsmObj::Node(ref node) => {
                                                       Some(mimir::CoordWrapper(geo::Coordinate {
                                                           x: node.lat,
                                                           y: node.lon,
                                                       }))
                                                   }
                                                   _ => None,
                                               }
                                           })
                                       });
            let (admin_id, insee_id) = match relation.tags.get("ref:INSEE") {
                Some(val) => (format!("admin:fr:{}", val.trim_left_matches('0')), val.trim_left_matches('0')),
                None => (format!("admin:osm:{}", relation.id), ""),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let boundary = mimirsbrunn::boundaries::build_boundary(&relation, &objects);
            let admin = mimir::Admin {
                id: admin_id,
                insee: insee_id.to_string(),
                level: level,
                label: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: Cell::new(0),
                coord: coord_centre,
                boundary: boundary,
            };
            administrative_regions.push(Rc::new(admin));
        }
    }
    return administrative_regions;
}

fn make_admin_geofinder(admins: AdminsVec) -> AdminGeoFinder {
    let mut geofinder = AdminGeoFinder::new();

    for a in admins {
        geofinder.add_admin(a);
    }
    geofinder
}

fn get_street_admin(admins_geofinder: &AdminGeoFinder,
                    obj_map: &HashMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>,
                    way: &osmpbfreader::objects::Way)
                    -> Vec<Rc<mimir::Admin>> {
    // for the moment we consider that the coord of the way is the coord of it's first node
    let coord = way.nodes
                   .iter()
                   .filter_map(|node_id| obj_map.get(&osmpbfreader::OsmId::Node(*node_id)))
                   .filter_map(|node_obj| {
                       if let &osmpbfreader::OsmObj::Node(ref node) = node_obj {
                           Some(geo::Coordinate {
                               x: node.lat,
                               y: node.lon,
                           })
                       } else {
                           None
                       }
                   })
                   .next();
    coord.map_or(vec![], |c| {
        admins_geofinder.get_admins_for_coord(&c)
    .iter()
    .cloned()
    .collect()
    })
}

fn streets(pbf: &mut OsmPbfReader, admins: AdminsVec) -> StreetsVec {
    let admins_geofinder = make_admin_geofinder(admins);

    let is_valid_obj = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Way(ref way) => {
                way.tags.get("highway").map_or(false, |x| !x.is_empty()) &&
                way.tags.get("name").map_or(false, |x| !x.is_empty())
            }
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("type").map_or(false, |v| v == "associatedStreet")
            }
            _ => false,
        }
    };

    let mut objs_map = osmpbfreader::get_objs_and_deps(pbf, is_valid_obj).unwrap();
    // Sometimes, streets can be devided into several "way"s that still have the same street name.
    // The reason why a street is devided may be that a part of the street become a bridge/tunne/etc.
    // In this case, a "relation" tagged with (type = associatedStreet) is used to group all these "way"s.
    // In order not to have doublons in autocompleion, we should keep only one "way" in the relation
    // and remove all the rest way whose street name is the same.
    let mut objs_to_remove = Vec::<osmpbfreader::OsmId>::new();

    for (_, rel_obj) in &objs_map {
        if let &osmpbfreader::OsmObj::Relation(ref rel) = rel_obj {
            let mut found = false;
            for ref_obj in &rel.refs {
                if let &osmpbfreader::OsmId::Way(_) = &ref_obj.member {
                    if !found {
                        found = true;
                        continue;
                    }
                    objs_to_remove.push(ref_obj.member.clone());
                };
            }
        };
    }

    for osm_id in objs_to_remove {
        objs_map.remove(&osm_id);
    }

    objs_map.iter()
            .filter_map(|(_, obj)| {
                if let &osmpbfreader::OsmObj::Way(ref way) = obj {
                    let admin = get_street_admin(&admins_geofinder, &objs_map, way);
                    way.tags.get("name").and_then(|way_name| {
                        Some(mimir::Street {
                            id: way.id.to_string(),
                            street_name: way_name.to_string(),
                            label: way_name.to_string(),
                            weight: 1,
                            administrative_regions: admin,
                        })
                    })
                } else {
                    None
                }
            })
            .collect()

}

fn main() {
    mimir::logger_init().unwrap();
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    debug!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);

    debug!("importing adminstrative region into Mimir");
    let admins = administrative_regions(&mut parsed_pbf, levels);

    let mut streets = streets(&mut parsed_pbf, admins.clone());

    for st in &mut streets {
    	for admin in &mut st.administrative_regions {
    		if admin.level == 8 {
    			admin.weight.set(admin.weight.get() + 1)
    		}
    	}
    }

    for st in &mut streets {
    	for admin in &mut st.administrative_regions {
    		if admin.level == 8 {
    			st.weight = admin.weight.get();
    			break;
    		}
    	}
    }

    if args.flag_import_way {
        debug!("importing streets into Mimir");
        let nb_streets = rubber.index("way",
                                      &args.flag_dataset,
                                      streets.into_iter())
                               .unwrap();
        info!("Nb of indexed street: {}", nb_streets);
    }

    let nb_admins = rubber.index("admin", &args.flag_dataset, admins.iter())
                          .unwrap();
    info!("Nb of indexed admin: {}", nb_admins);

}
<commit_msg>mimir: parameter cety-level added to calculate weight<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

#[macro_use]
extern crate log;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimir;
extern crate rs_es;
extern crate chrono;

extern crate mimirsbrunn;
extern crate serde;
extern crate geo;

use std::collections::{HashSet, HashMap};
use mimir::rubber::Rubber;

pub type AdminsVec = Vec<Rc<mimir::Admin>>;
pub type StreetsVec = Vec<mimir::Street>;
pub type OsmPbfReader = osmpbfreader::OsmPbfReader<std::fs::File>;

use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use std::rc::Rc;
use std::cell::Cell;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_city_level: u32,
    flag_connection_string: String,
    flag_import_way: bool,
    flag_import_admin: bool,
    flag_dataset: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] [--import-way] [--import-admin] [--dataset=<dataset>] [--city-level=<level>] --level=<level> ...

Options:
    -h, --help              Show this message.
    -i, --input=<file>      OSM PBF file.
    -l, --level=<level>     Admin levels to keep.
    -C, --city-level=<level>
                            City level to calculate weight, [default: 8]
    -w, --import-way        Import ways
    -a, --import-admin      Import admins
    -c, --connection-string=<connection-string>
                            Elasticsearch parameters, [default: http://localhost:9200/munin]
    -d, --dataset=<dataset>
                            Name of the dataset, [default: fr]

";

#[derive(Debug)]
struct AdminMatcher {
    admin_levels: HashSet<u32>,
}
impl AdminMatcher {
    pub fn new(levels: HashSet<u32>) -> AdminMatcher {
        AdminMatcher { admin_levels: levels }
    }

    pub fn is_admin(&self, obj: &osmpbfreader::OsmObj) -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("boundary").map_or(false, |v| v == "administrative") &&
                rel.tags.get("admin_level").map_or(false, |lvl| {
                    self.admin_levels.contains(&lvl.parse::<u32>().unwrap_or(0))
                })
            }
            _ => false,
        }
    }
}

fn parse_osm_pbf(path: &str) -> OsmPbfReader {
    let path = std::path::Path::new(&path);
    osmpbfreader::OsmPbfReader::new(std::fs::File::open(&path).unwrap())
}

fn administrative_regions(pbf: &mut OsmPbfReader, levels: HashSet<u32>) -> AdminsVec {
    let mut administrative_regions = AdminsVec::new();
    let matcher = AdminMatcher::new(levels);
    let objects = osmpbfreader::get_objs_and_deps(pbf, |o| matcher.is_admin(o)).unwrap();
    // load administratives regions
    for (_, obj) in &objects {
        if !matcher.is_admin(&obj) {
            continue;
        }
        if let &osmpbfreader::OsmObj::Relation(ref relation) = obj {
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    warn!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            // admininstrative region without coordinates
            let coord_centre = relation.refs
                                       .iter()
                                       .find(|rf| rf.role == "admin_centre")
                                       .and_then(|r| {
                                           objects.get(&r.member).and_then(|value| {
                                               match value {
                                                   &osmpbfreader::OsmObj::Node(ref node) => {
                                                       Some(mimir::CoordWrapper(geo::Coordinate {
                                                           x: node.lat,
                                                           y: node.lon,
                                                       }))
                                                   }
                                                   _ => None,
                                               }
                                           })
                                       });
            let (admin_id, insee_id) = match relation.tags.get("ref:INSEE") {
                Some(val) => (format!("admin:fr:{}", val.trim_left_matches('0')), val.trim_left_matches('0')),
                None => (format!("admin:osm:{}", relation.id), ""),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let boundary = mimirsbrunn::boundaries::build_boundary(&relation, &objects);
            let admin = mimir::Admin {
                id: admin_id,
                insee: insee_id.to_string(),
                level: level,
                label: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: Cell::new(0),
                coord: coord_centre,
                boundary: boundary,
            };
            administrative_regions.push(Rc::new(admin));
        }
    }
    return administrative_regions;
}

fn make_admin_geofinder(admins: AdminsVec) -> AdminGeoFinder {
    let mut geofinder = AdminGeoFinder::new();

    for a in admins {
        geofinder.add_admin(a);
    }
    geofinder
}

fn get_street_admin(admins_geofinder: &AdminGeoFinder,
                    obj_map: &HashMap<osmpbfreader::OsmId, osmpbfreader::OsmObj>,
                    way: &osmpbfreader::objects::Way)
                    -> Vec<Rc<mimir::Admin>> {
    // for the moment we consider that the coord of the way is the coord of it's first node
    let coord = way.nodes
                   .iter()
                   .filter_map(|node_id| obj_map.get(&osmpbfreader::OsmId::Node(*node_id)))
                   .filter_map(|node_obj| {
                       if let &osmpbfreader::OsmObj::Node(ref node) = node_obj {
                           Some(geo::Coordinate {
                               x: node.lat,
                               y: node.lon,
                           })
                       } else {
                           None
                       }
                   })
                   .next();
    coord.map_or(vec![], |c| {
        admins_geofinder.get_admins_for_coord(&c)
    .iter()
    .cloned()
    .collect()
    })
}

fn streets(pbf: &mut OsmPbfReader, admins: AdminsVec) -> StreetsVec {
    let admins_geofinder = make_admin_geofinder(admins);

    let is_valid_obj = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Way(ref way) => {
                way.tags.get("highway").map_or(false, |x| !x.is_empty()) &&
                way.tags.get("name").map_or(false, |x| !x.is_empty())
            }
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags.get("type").map_or(false, |v| v == "associatedStreet")
            }
            _ => false,
        }
    };

    let mut objs_map = osmpbfreader::get_objs_and_deps(pbf, is_valid_obj).unwrap();
    // Sometimes, streets can be devided into several "way"s that still have the same street name.
    // The reason why a street is devided may be that a part of the street become a bridge/tunne/etc.
    // In this case, a "relation" tagged with (type = associatedStreet) is used to group all these "way"s.
    // In order not to have doublons in autocompleion, we should keep only one "way" in the relation
    // and remove all the rest way whose street name is the same.
    let mut objs_to_remove = Vec::<osmpbfreader::OsmId>::new();

    for (_, rel_obj) in &objs_map {
        if let &osmpbfreader::OsmObj::Relation(ref rel) = rel_obj {
            let mut found = false;
            for ref_obj in &rel.refs {
                if let &osmpbfreader::OsmId::Way(_) = &ref_obj.member {
                    if !found {
                        found = true;
                        continue;
                    }
                    objs_to_remove.push(ref_obj.member.clone());
                };
            }
        };
    }

    for osm_id in objs_to_remove {
        objs_map.remove(&osm_id);
    }

    objs_map.iter()
            .filter_map(|(_, obj)| {
                if let &osmpbfreader::OsmObj::Way(ref way) = obj {
                    let admin = get_street_admin(&admins_geofinder, &objs_map, way);
                    way.tags.get("name").and_then(|way_name| {
                        Some(mimir::Street {
                            id: way.id.to_string(),
                            street_name: way_name.to_string(),
                            label: way_name.to_string(),
                            weight: 1,
                            administrative_regions: admin,
                        })
                    })
                } else {
                    None
                }
            })
            .collect()

}

fn main() {
    mimir::logger_init().unwrap();
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let city_level = args.flag_city_level;
    let mut parsed_pbf = parse_osm_pbf(&args.flag_input);
    debug!("creation of indexes");
    let mut rubber = Rubber::new(&args.flag_connection_string);

    debug!("importing adminstrative region into Mimir");
    let admins = administrative_regions(&mut parsed_pbf, levels);

    let mut streets = streets(&mut parsed_pbf, admins.clone());

    for st in &mut streets {
    	for admin in &mut st.administrative_regions {
    		if admin.level == city_level {
    			admin.weight.set(admin.weight.get() + 1)
    		}
    	}
    }

    for st in &mut streets {
    	for admin in &mut st.administrative_regions {
    		if admin.level == city_level {
    			st.weight = admin.weight.get();
    			break;
    		}
    	}
    }

    if args.flag_import_way {
        debug!("importing streets into Mimir");
        let nb_streets = rubber.index("way",
                                      &args.flag_dataset,
                                      streets.into_iter())
                               .unwrap();
        info!("Nb of indexed street: {}", nb_streets);
    }

    let nb_admins = rubber.index("admin", &args.flag_dataset, admins.iter())
                          .unwrap();
    info!("Nb of indexed admin: {}", nb_admins);

}
<|endoftext|>


<commit_before>// Copyright 2016 Timothée Ravier
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use libc::*;
use std::ops::Drop;

use ffi::*;
use handle::Handle;
use error::Error;

pub struct Message<'a> {
    msg: bus1_msg,

    peer_fd: c_int,
    handles: &'a [u64],

    // type: BUS1_MSG,
    // destination: u64,
    // uid: uid_t,
    // gid: gid_t,
    // pid: pid_t,
    // tid: tid_t,
    // struct iovecs *vecs,
    // size_t n_vecs,
    // B1Handle **handles,
    // size_t n_handles,
    // int *fds,
    // size_t n_fds
}

impl<'a> Message<'a> {
    pub fn new(msg: bus1_msg, fd: c_int, handles: &'a [u64],) -> Message<'a> {
        debug!("New message received on handle: {}, offset: {}, n_bytes: {}, n_handles: {}, n_fds: {}",
               msg.destination, msg.offset, msg.n_bytes, msg.n_handles, msg.n_fds);
        Message {
            msg: msg,
            peer_fd: fd,
            handles: handles,
        }
    }

    pub fn type_(&self) -> u64 {
        self.msg.type_
    }

    pub fn flags(&self) -> u64 {
        self.msg.flags
    }

    pub fn destination(&self) -> Handle {
        Handle::new(self.msg.destination)
    }

    pub fn release_slice(self) {
        ()
    }

    fn _release_slice(&self) -> Option<Error> {
        unsafe {
            // FIXME: bus1 ioctl is readwrite but does not use this arguement
            let mut fixme_offset = self.msg.offset;
            let r = bus1_cmd_slice_release(self.peer_fd, &mut fixme_offset);
            if r < 0 {
                debug!("Could not drop slice (offset: {}, n_bytes: {})", self.msg.offset, self.msg.n_bytes);
                return Some(Error::last_os_error());
            }
        }
        debug!("Dropped slice (offset: {}, n_bytes: {})", self.msg.offset, self.msg.n_bytes);
        None
    }

    pub fn get_handles(&self) -> &'a [u64] {
        self.handles
    }
}

impl<'a> Drop for Message<'a> {
    fn drop(&mut self) {
        match self._release_slice() {
            Some(e) => panic!(e), // FIXME: pragma/feature debug
            None => ()
        }
    }
}
<commit_msg>message: add derive(Debug) to Message struct<commit_after>// Copyright 2016 Timothée Ravier
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use libc::*;
use std::ops::Drop;

use ffi::*;
use handle::Handle;
use error::Error;

#[derive(Debug)]
pub struct Message<'a> {
    msg: bus1_msg,

    peer_fd: c_int,
    handles: &'a [u64],

    // type: BUS1_MSG,
    // destination: u64,
    // uid: uid_t,
    // gid: gid_t,
    // pid: pid_t,
    // tid: tid_t,
    // struct iovecs *vecs,
    // size_t n_vecs,
    // B1Handle **handles,
    // size_t n_handles,
    // int *fds,
    // size_t n_fds
}

impl<'a> Message<'a> {
    pub fn new(msg: bus1_msg, fd: c_int, handles: &'a [u64],) -> Message<'a> {
        debug!("New message received on handle: {}, offset: {}, n_bytes: {}, n_handles: {}, n_fds: {}",
               msg.destination, msg.offset, msg.n_bytes, msg.n_handles, msg.n_fds);
        Message {
            msg: msg,
            peer_fd: fd,
            handles: handles,
        }
    }

    pub fn type_(&self) -> u64 {
        self.msg.type_
    }

    pub fn flags(&self) -> u64 {
        self.msg.flags
    }

    pub fn destination(&self) -> Handle {
        Handle::new(self.msg.destination)
    }

    pub fn release_slice(self) {
        ()
    }

    fn _release_slice(&self) -> Option<Error> {
        unsafe {
            // FIXME: bus1 ioctl is readwrite but does not use this arguement
            let mut fixme_offset = self.msg.offset;
            let r = bus1_cmd_slice_release(self.peer_fd, &mut fixme_offset);
            if r < 0 {
                debug!("Could not drop slice (offset: {}, n_bytes: {})", self.msg.offset, self.msg.n_bytes);
                return Some(Error::last_os_error());
            }
        }
        debug!("Dropped slice (offset: {}, n_bytes: {})", self.msg.offset, self.msg.n_bytes);
        None
    }

    pub fn get_handles(&self) -> &'a [u64] {
        self.handles
    }
}

impl<'a> Drop for Message<'a> {
    fn drop(&mut self) {
        match self._release_slice() {
            Some(e) => panic!(e), // FIXME: pragma/feature debug
            None => ()
        }
    }
}
<|endoftext|>


<commit_before>//! Store and manipulate Kafka messages.
use rdsys;
use rdsys::types::*;

use std::ffi::CStr;
use std::fmt;
use std::marker::PhantomData;
use std::slice;
use std::str;

use consumer::{Consumer, ConsumerContext};

/// Timestamp of a message
#[derive(Debug,PartialEq,Eq,Clone,Copy)]
pub enum Timestamp {
    NotAvailable,
    CreateTime(i64),
    LogAppendTime(i64)
}

impl Timestamp {
    pub fn to_millis(&self) -> Option<i64> {
        match *self {
            Timestamp::NotAvailable => None,
            Timestamp::CreateTime(-1) | Timestamp::LogAppendTime(-1) => None,
            Timestamp::CreateTime(t) | Timestamp::LogAppendTime(t) => Some(t),
        }
    }
}

/// The `Message` trait provides access to the fields of a generic Kafka message.
pub trait Message {
    /// Returns the key of the message, or None if there is no key.
    fn key(&self) -> Option<&[u8]>;

    /// Returns the payload of the message, or None if there is no payload.
    fn payload(&self) -> Option<&[u8]>;

    /// Returns the source topic of the message.
    fn topic(&self) -> &str;

    /// Returns the partition number where the message is stored.
    fn partition(&self) -> i32;

    /// Returns the offset of the message.
    fn offset(&self) -> i64;

    /// Returns the message timestamp for a consumed message if available.
    fn timestamp(&self) -> Timestamp;

    /// Converts the raw bytes of the payload to a reference of the specified type, that points to the
    /// same data inside the message and without performing any memory allocation
    fn payload_view<P: ?Sized + FromBytes>(&self) -> Option<Result<&P, P::Error>> {
        self.payload().map(P::from_bytes)
    }

    /// Converts the raw bytes of the key to a reference of the specified type, that points to the
    /// same data inside the message and without performing any memory allocation
    fn key_view<K: ?Sized + FromBytes>(&self) -> Option<Result<&K, K::Error>> {
        self.key().map(K::from_bytes)
    }
}

/// A native librdkafka message. The content of the message is stored in the receiving buffer of
/// the consumer, to avoid memory allocations. As such, `BorrowedMessage` cannot outlive the
/// consumer it is received from.
/// `BorrowedMessage`s are removed from the consumer buffer once they are dropped. Holding
/// references to many messages might cause the memory of the consumer to fill up and stop.
/// To transform a `BorrowedMessage` into a `OwnedMessage`, use the `detach` method.
pub struct BorrowedMessage<'a> {
    ptr: *mut RDKafkaMessage,
    _p: PhantomData<&'a u8>,
}

impl<'a> fmt::Debug for BorrowedMessage<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Message {{ ptr: {:?} }}", self.ptr())
    }
}

impl<'a> BorrowedMessage<'a> {
    /// Creates a new `BorrowedMessage` that wraps the native Kafka message pointer. The lifetime of the
    /// message will be bound to the lifetime of the consumer passed as parameter.
    pub fn new<C, X>(ptr: *mut RDKafkaMessage, _consumer: &'a C) -> BorrowedMessage<'a>
        where X: ConsumerContext,
              C: Consumer<X> {
        BorrowedMessage {
            ptr: ptr,
            _p: PhantomData,
        }
    }

    /// Returns a pointer to the RDKafkaMessage.
    pub fn ptr(&self) -> *mut RDKafkaMessage {
        self.ptr
    }

    /// Returns a pointer to the message's RDKafkaTopic
    pub fn topic_ptr(&self) -> *mut RDKafkaTopic {
        unsafe { (*self.ptr).rkt }
    }

    /// Returns the length of the key field of the message.
    pub fn key_len(&self) -> usize {
        unsafe { (*self.ptr).key_len }
    }

    /// Returns the length of the payload field of the message.
    pub fn payload_len(&self) -> usize {
        unsafe { (*self.ptr).len }
    }

    /// Clones the content of the `BorrowedMessage` and returns an `OwnedMessage`, that can
    /// outlive the consumer. This operation requires memory allocation and can be expensive.
    pub fn detach(&self) -> OwnedMessage {
        OwnedMessage {
            key: self.key().map(|k| k.to_vec()),
            payload: self.payload().map(|p| p.to_vec()),
            topic: self.topic().to_owned(),
            timestamp: self.timestamp(),
            partition: self.partition(),
            offset: self.offset(),
        }
    }
}

impl<'a> Message for BorrowedMessage<'a> {
    fn key(&self) -> Option<&[u8]> {
        unsafe {
            if (*self.ptr).key.is_null() {
                None
            } else {
                Some(slice::from_raw_parts::<u8>((*self.ptr).key as *const u8, (*self.ptr).key_len))
            }
        }
    }

    fn payload(&self) -> Option<&[u8]> {
        unsafe {
            if (*self.ptr).payload.is_null() {
                None
            } else {
                Some(slice::from_raw_parts::<u8>((*self.ptr).payload as *const u8, (*self.ptr).len))
            }
        }
    }

    fn topic(&self) -> &str {
         unsafe {
             CStr::from_ptr(rdsys::rd_kafka_topic_name((*self.ptr).rkt))
                 .to_str()
                 .expect("Topic name is not valid UTF-8")
         }
     }

    fn partition(&self) -> i32 {
        unsafe { (*self.ptr).partition }
    }

    fn offset(&self) -> i64 {
        unsafe { (*self.ptr).offset }
    }

    // TODO: -1 should be "Not available"
    fn timestamp(&self) -> Timestamp {
        let mut timestamp_type = rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_NOT_AVAILABLE;
        let timestamp = unsafe {
            rdsys::rd_kafka_message_timestamp(
                self.ptr,
                &mut timestamp_type
            )

        };
        match timestamp_type {
            rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_NOT_AVAILABLE => Timestamp::NotAvailable,
            rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_CREATE_TIME => Timestamp::CreateTime(timestamp),
            rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_LOG_APPEND_TIME => Timestamp::LogAppendTime(timestamp)
        }
    }
}

impl<'a> Drop for BorrowedMessage<'a> {
    fn drop(&mut self) {
        trace!("Destroying message {:?}", self);
        unsafe { rdsys::rd_kafka_message_destroy(self.ptr) };
    }
}

//
// ********** OWNED MESSAGE **********
//

/// An `OwnedMessage` can be created from a `BorrowedMessage` using the `detach` method.
/// `OwnedMessage`s don't hold any reference to the consumer, and don't use any memory inside the
/// consumer buffer.
pub struct OwnedMessage {
    payload: Option<Vec<u8>>,
    key: Option<Vec<u8>>,
    topic: String,
    timestamp: Timestamp,
    partition: i32,
    offset: i64
}

impl Message for OwnedMessage {
    fn key(&self) -> Option<&[u8]> {
        match self.key {
            Some(ref k) => Some(k.as_slice()),
            None => None,
        }
    }

    fn payload(&self) -> Option<&[u8]> {
        match self.payload {
            Some(ref p) => Some(p.as_slice()),
            None => None,
        }
    }

    fn topic(&self) -> &str {
        self.topic.as_ref()
    }

    fn partition(&self) -> i32 {
        self.partition
    }

    fn offset(&self) -> i64 {
        self.offset
    }

    fn timestamp(&self) -> Timestamp {
        self.timestamp
    }
}

/// Given a reference to a byte array, returns a different view of the same data.
/// No copy of the data should be performed.
pub trait FromBytes {
    type Error;
    fn from_bytes(&[u8]) -> Result<&Self, Self::Error>;
}

impl FromBytes for [u8] {
    type Error = ();
    fn from_bytes(bytes: &[u8]) -> Result<&Self, Self::Error> {
        Ok(bytes)
    }
}

impl FromBytes for str {
    type Error = str::Utf8Error;
    fn from_bytes(bytes: &[u8]) -> Result<&Self, Self::Error> {
        str::from_utf8(bytes)
    }
}

/// Given some data, returns the byte representation of that data.
/// No copy of the data should be performed.
pub trait ToBytes {
    fn to_bytes(&self) -> &[u8];
}

impl ToBytes for [u8] {
    fn to_bytes(&self) -> &[u8] {
        self
    }
}

impl ToBytes for str {
    fn to_bytes(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToBytes for Vec<u8> {
    fn to_bytes(&self) -> &[u8] {
        self.as_slice()
    }
}

impl ToBytes for String {
    fn to_bytes(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl<'a, T: ToBytes> ToBytes for &'a T {
    fn to_bytes(&self) -> &[u8] {
        (*self).to_bytes()
    }
}

impl ToBytes for () {
    fn to_bytes(&self) -> &[u8] {
        &[]
    }
}
<commit_msg>Add way to create an OwnedMessage<commit_after>//! Store and manipulate Kafka messages.
use rdsys;
use rdsys::types::*;

use std::ffi::CStr;
use std::fmt;
use std::marker::PhantomData;
use std::slice;
use std::str;

use consumer::{Consumer, ConsumerContext};

/// Timestamp of a message
#[derive(Debug,PartialEq,Eq,Clone,Copy)]
pub enum Timestamp {
    NotAvailable,
    CreateTime(i64),
    LogAppendTime(i64)
}

impl Timestamp {
    pub fn to_millis(&self) -> Option<i64> {
        match *self {
            Timestamp::NotAvailable => None,
            Timestamp::CreateTime(-1) | Timestamp::LogAppendTime(-1) => None,
            Timestamp::CreateTime(t) | Timestamp::LogAppendTime(t) => Some(t),
        }
    }
}

/// The `Message` trait provides access to the fields of a generic Kafka message.
pub trait Message {
    /// Returns the key of the message, or None if there is no key.
    fn key(&self) -> Option<&[u8]>;

    /// Returns the payload of the message, or None if there is no payload.
    fn payload(&self) -> Option<&[u8]>;

    /// Returns the source topic of the message.
    fn topic(&self) -> &str;

    /// Returns the partition number where the message is stored.
    fn partition(&self) -> i32;

    /// Returns the offset of the message.
    fn offset(&self) -> i64;

    /// Returns the message timestamp for a consumed message if available.
    fn timestamp(&self) -> Timestamp;

    /// Converts the raw bytes of the payload to a reference of the specified type, that points to the
    /// same data inside the message and without performing any memory allocation
    fn payload_view<P: ?Sized + FromBytes>(&self) -> Option<Result<&P, P::Error>> {
        self.payload().map(P::from_bytes)
    }

    /// Converts the raw bytes of the key to a reference of the specified type, that points to the
    /// same data inside the message and without performing any memory allocation
    fn key_view<K: ?Sized + FromBytes>(&self) -> Option<Result<&K, K::Error>> {
        self.key().map(K::from_bytes)
    }
}

/// A native librdkafka message. The content of the message is stored in the receiving buffer of
/// the consumer, to avoid memory allocations. As such, `BorrowedMessage` cannot outlive the
/// consumer it is received from.
/// `BorrowedMessage`s are removed from the consumer buffer once they are dropped. Holding
/// references to many messages might cause the memory of the consumer to fill up and stop.
/// To transform a `BorrowedMessage` into a `OwnedMessage`, use the `detach` method.
pub struct BorrowedMessage<'a> {
    ptr: *mut RDKafkaMessage,
    _p: PhantomData<&'a u8>,
}

impl<'a> fmt::Debug for BorrowedMessage<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Message {{ ptr: {:?} }}", self.ptr())
    }
}

impl<'a> BorrowedMessage<'a> {
    /// Creates a new `BorrowedMessage` that wraps the native Kafka message pointer. The lifetime of the
    /// message will be bound to the lifetime of the consumer passed as parameter.
    pub fn new<C, X>(ptr: *mut RDKafkaMessage, _consumer: &'a C) -> BorrowedMessage<'a>
        where X: ConsumerContext,
              C: Consumer<X> {
        BorrowedMessage {
            ptr: ptr,
            _p: PhantomData,
        }
    }

    /// Returns a pointer to the RDKafkaMessage.
    pub fn ptr(&self) -> *mut RDKafkaMessage {
        self.ptr
    }

    /// Returns a pointer to the message's RDKafkaTopic
    pub fn topic_ptr(&self) -> *mut RDKafkaTopic {
        unsafe { (*self.ptr).rkt }
    }

    /// Returns the length of the key field of the message.
    pub fn key_len(&self) -> usize {
        unsafe { (*self.ptr).key_len }
    }

    /// Returns the length of the payload field of the message.
    pub fn payload_len(&self) -> usize {
        unsafe { (*self.ptr).len }
    }

    /// Clones the content of the `BorrowedMessage` and returns an `OwnedMessage`, that can
    /// outlive the consumer. This operation requires memory allocation and can be expensive.
    pub fn detach(&self) -> OwnedMessage {
        OwnedMessage {
            key: self.key().map(|k| k.to_vec()),
            payload: self.payload().map(|p| p.to_vec()),
            topic: self.topic().to_owned(),
            timestamp: self.timestamp(),
            partition: self.partition(),
            offset: self.offset(),
        }
    }
}

impl<'a> Message for BorrowedMessage<'a> {
    fn key(&self) -> Option<&[u8]> {
        unsafe {
            if (*self.ptr).key.is_null() {
                None
            } else {
                Some(slice::from_raw_parts::<u8>((*self.ptr).key as *const u8, (*self.ptr).key_len))
            }
        }
    }

    fn payload(&self) -> Option<&[u8]> {
        unsafe {
            if (*self.ptr).payload.is_null() {
                None
            } else {
                Some(slice::from_raw_parts::<u8>((*self.ptr).payload as *const u8, (*self.ptr).len))
            }
        }
    }

    fn topic(&self) -> &str {
         unsafe {
             CStr::from_ptr(rdsys::rd_kafka_topic_name((*self.ptr).rkt))
                 .to_str()
                 .expect("Topic name is not valid UTF-8")
         }
     }

    fn partition(&self) -> i32 {
        unsafe { (*self.ptr).partition }
    }

    fn offset(&self) -> i64 {
        unsafe { (*self.ptr).offset }
    }

    // TODO: -1 should be "Not available"
    fn timestamp(&self) -> Timestamp {
        let mut timestamp_type = rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_NOT_AVAILABLE;
        let timestamp = unsafe {
            rdsys::rd_kafka_message_timestamp(
                self.ptr,
                &mut timestamp_type
            )

        };
        match timestamp_type {
            rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_NOT_AVAILABLE => Timestamp::NotAvailable,
            rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_CREATE_TIME => Timestamp::CreateTime(timestamp),
            rdsys::rd_kafka_timestamp_type_t::RD_KAFKA_TIMESTAMP_LOG_APPEND_TIME => Timestamp::LogAppendTime(timestamp)
        }
    }
}

impl<'a> Drop for BorrowedMessage<'a> {
    fn drop(&mut self) {
        trace!("Destroying message {:?}", self);
        unsafe { rdsys::rd_kafka_message_destroy(self.ptr) };
    }
}

//
// ********** OWNED MESSAGE **********
//

/// An `OwnedMessage` can be created from a `BorrowedMessage` using the `detach` method.
/// `OwnedMessage`s don't hold any reference to the consumer, and don't use any memory inside the
/// consumer buffer.
pub struct OwnedMessage {
    payload: Option<Vec<u8>>,
    key: Option<Vec<u8>>,
    topic: String,
    timestamp: Timestamp,
    partition: i32,
    offset: i64
}

impl OwnedMessage {
    /// Create a new message with the specified content. Mainly useful for writing tests.
    pub fn new(
        payload: Option<Vec<u8>>,
        key: Option<Vec<u8>>,
        topic: String,
        timestamp: Timestamp,
        partition: i32,
        offset: i64
    ) -> OwnedMessage {
        OwnedMessage {
            payload: payload,
            key: key,
            topic: topic,
            timestamp: timestamp,
            partition: partition,
            offset: offset
        }
    }
}

impl Message for OwnedMessage {
    fn key(&self) -> Option<&[u8]> {
        match self.key {
            Some(ref k) => Some(k.as_slice()),
            None => None,
        }
    }

    fn payload(&self) -> Option<&[u8]> {
        match self.payload {
            Some(ref p) => Some(p.as_slice()),
            None => None,
        }
    }

    fn topic(&self) -> &str {
        self.topic.as_ref()
    }

    fn partition(&self) -> i32 {
        self.partition
    }

    fn offset(&self) -> i64 {
        self.offset
    }

    fn timestamp(&self) -> Timestamp {
        self.timestamp
    }
}

/// Given a reference to a byte array, returns a different view of the same data.
/// No copy of the data should be performed.
pub trait FromBytes {
    type Error;
    fn from_bytes(&[u8]) -> Result<&Self, Self::Error>;
}

impl FromBytes for [u8] {
    type Error = ();
    fn from_bytes(bytes: &[u8]) -> Result<&Self, Self::Error> {
        Ok(bytes)
    }
}

impl FromBytes for str {
    type Error = str::Utf8Error;
    fn from_bytes(bytes: &[u8]) -> Result<&Self, Self::Error> {
        str::from_utf8(bytes)
    }
}

/// Given some data, returns the byte representation of that data.
/// No copy of the data should be performed.
pub trait ToBytes {
    fn to_bytes(&self) -> &[u8];
}

impl ToBytes for [u8] {
    fn to_bytes(&self) -> &[u8] {
        self
    }
}

impl ToBytes for str {
    fn to_bytes(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToBytes for Vec<u8> {
    fn to_bytes(&self) -> &[u8] {
        self.as_slice()
    }
}

impl ToBytes for String {
    fn to_bytes(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl<'a, T: ToBytes> ToBytes for &'a T {
    fn to_bytes(&self) -> &[u8] {
        (*self).to_bytes()
    }
}

impl ToBytes for () {
    fn to_bytes(&self) -> &[u8] {
        &[]
    }
}
<|endoftext|>


<commit_before>//! This crate provides the `#[quickcheck]` attribute. Its use is
//! documented in the `quickcheck` crate.

#![crate_name = "quickcheck_macros"]
#![crate_type = "dylib"]
#![doc(html_root_url = "http://burntsushi.net/rustdoc/quickcheck")]

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc_plugin;

use syntax::abi;
use syntax::ast;
use syntax::ast::Ty_::TyBareFn;
use syntax::ast_util;
use syntax::codemap;
use syntax::parse::token;
use syntax::ext::base::{ExtCtxt, MultiModifier, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use rustc_plugin::Registry;

/// For the `#[quickcheck]` attribute. Do not use.
#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(token::intern("quickcheck"),
                                  MultiModifier(Box::new(expand_meta_quickcheck)));
}

/// Expands the `#[quickcheck]` attribute.
///
/// Expands:
/// ```
/// #[quickcheck]
/// fn check_something(_: usize) -> bool {
///     true
/// }
/// ```
/// to:
/// ```
/// #[test]
/// fn check_something() {
///     fn check_something(_: usize) -> bool {
///         true
///     }
///     ::quickcheck::quickcheck(check_something as fn(usize) -> bool)
/// }
/// ```
fn expand_meta_quickcheck(cx: &mut ExtCtxt,
                          span: codemap::Span,
                          _: &ast::MetaItem,
                          annot_item: Annotatable) -> Annotatable {
    let item = annot_item.expect_item();
    match item.node {
        ast::ItemFn(ref decl, unsafety, _constness, abi, _, _) => {
            let prop_ident = cx.expr_ident(span, item.ident);
            let prop_ty = cx.ty(span, TyBareFn(P(ast::BareFnTy {
                unsafety: unsafety,
                abi: abi,
                lifetimes: vec![],
                decl: decl.clone(),
            })));
            let inner_ident = cx.expr_cast(span, prop_ident, prop_ty);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        ast::ItemStatic(..) => {
            let inner_ident = cx.expr_ident(span, item.ident);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        _ => {
            cx.span_err(
                span, "#[quickcheck] only supported on statics and functions");
        }
    }
    Annotatable::Item(item)
}

fn wrap_item(cx: &mut ExtCtxt,
             span: codemap::Span,
             item: &ast::Item,
             inner_ident: P<ast::Expr>) -> Annotatable {
    // Copy original function without attributes
    let prop = P(ast::Item {attrs: Vec::new(), ..item.clone()});
    // ::quickcheck::quickcheck
    let check_ident = token::str_to_ident("quickcheck");
    let check_path = vec!(check_ident, check_ident);
    // Wrap original function in new outer function,
    // calling ::quickcheck::quickcheck()
    let fn_decl = P(codemap::respan(span, ast::DeclItem(prop.clone())));
    let inner_fn =
        P(codemap::respan(span, ast::StmtDecl(fn_decl, ast::DUMMY_NODE_ID)));
    let check_call = cx.expr_call_global(span, check_path, vec![inner_ident]);
    let body = cx.block(span, vec![inner_fn], Some(check_call));
    let test = item_fn(cx, span, item, body);

    // Copy attributes from original function
    let mut attrs = item.attrs.clone();
    // Add #[test] attribute
    attrs.push(cx.attribute(
        span, cx.meta_word(span, token::intern_and_get_ident("test"))));
    // Attach the attributes to the outer function
    Annotatable::Item(P(ast::Item {attrs: attrs, ..(*test).clone()}))
}

fn item_fn(cx: &mut ExtCtxt, span: codemap::Span,
           towrap_item: &ast::Item, body: P<ast::Block>) -> P<ast::Item> {
    let decl = P(ast::FnDecl {
        inputs: vec![],
        output: ast::FunctionRetTy::DefaultReturn(span),
        variadic: false,
    });
    let item = ast::ItemFn(decl,
                           ast::Unsafety::Normal,
                           ast::Constness::NotConst,
                           abi::Rust,
                           ast_util::empty_generics(),
                           body);
    cx.item(span, towrap_item.ident, vec![], item)
}
<commit_msg>Fixed regression in nightly<commit_after>//! This crate provides the `#[quickcheck]` attribute. Its use is
//! documented in the `quickcheck` crate.

#![crate_name = "quickcheck_macros"]
#![crate_type = "dylib"]
#![doc(html_root_url = "http://burntsushi.net/rustdoc/quickcheck")]

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc_plugin;

use syntax::abi;
use syntax::ast;
use syntax::ast::Ty_::TyBareFn;
use syntax::codemap;
use syntax::parse::token;
use syntax::ext::base::{ExtCtxt, MultiModifier, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use rustc_plugin::Registry;

/// For the `#[quickcheck]` attribute. Do not use.
#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(token::intern("quickcheck"),
                                  MultiModifier(Box::new(expand_meta_quickcheck)));
}

/// Expands the `#[quickcheck]` attribute.
///
/// Expands:
/// ```
/// #[quickcheck]
/// fn check_something(_: usize) -> bool {
///     true
/// }
/// ```
/// to:
/// ```
/// #[test]
/// fn check_something() {
///     fn check_something(_: usize) -> bool {
///         true
///     }
///     ::quickcheck::quickcheck(check_something as fn(usize) -> bool)
/// }
/// ```
fn expand_meta_quickcheck(cx: &mut ExtCtxt,
                          span: codemap::Span,
                          _: &ast::MetaItem,
                          annot_item: Annotatable) -> Annotatable {
    let item = annot_item.expect_item();
    match item.node {
        ast::ItemFn(ref decl, unsafety, _constness, abi, _, _) => {
            let prop_ident = cx.expr_ident(span, item.ident);
            let prop_ty = cx.ty(span, TyBareFn(P(ast::BareFnTy {
                unsafety: unsafety,
                abi: abi,
                lifetimes: vec![],
                decl: decl.clone(),
            })));
            let inner_ident = cx.expr_cast(span, prop_ident, prop_ty);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        ast::ItemStatic(..) => {
            let inner_ident = cx.expr_ident(span, item.ident);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        _ => {
            cx.span_err(
                span, "#[quickcheck] only supported on statics and functions");
        }
    }
    Annotatable::Item(item)
}

fn wrap_item(cx: &mut ExtCtxt,
             span: codemap::Span,
             item: &ast::Item,
             inner_ident: P<ast::Expr>) -> Annotatable {
    // Copy original function without attributes
    let prop = P(ast::Item {attrs: Vec::new(), ..item.clone()});
    // ::quickcheck::quickcheck
    let check_ident = token::str_to_ident("quickcheck");
    let check_path = vec!(check_ident, check_ident);
    // Wrap original function in new outer function,
    // calling ::quickcheck::quickcheck()
    let fn_decl = P(codemap::respan(span, ast::DeclItem(prop.clone())));
    let inner_fn =
        P(codemap::respan(span, ast::StmtDecl(fn_decl, ast::DUMMY_NODE_ID)));
    let check_call = cx.expr_call_global(span, check_path, vec![inner_ident]);
    let body = cx.block(span, vec![inner_fn], Some(check_call));
    let test = item_fn(cx, span, item, body);

    // Copy attributes from original function
    let mut attrs = item.attrs.clone();
    // Add #[test] attribute
    attrs.push(cx.attribute(
        span, cx.meta_word(span, token::intern_and_get_ident("test"))));
    // Attach the attributes to the outer function
    Annotatable::Item(P(ast::Item {attrs: attrs, ..(*test).clone()}))
}

fn item_fn(cx: &mut ExtCtxt, span: codemap::Span,
           towrap_item: &ast::Item, body: P<ast::Block>) -> P<ast::Item> {
    let decl = P(ast::FnDecl {
        inputs: vec![],
        output: ast::FunctionRetTy::DefaultReturn(span),
        variadic: false,
    });
    let item = ast::ItemFn(decl,
                           ast::Unsafety::Normal,
                           ast::Constness::NotConst,
                           abi::Rust,
                           ast::Generics::default(),
                           body);
    cx.item(span, towrap_item.ident, vec![], item)
}
<|endoftext|>