<commit_before>extern crate slack;

use std::collections::BTreeMap;
use std::env;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;
use std::sync::mpsc::channel;
use std::thread;

use rustc_serialize::json::{self, Json, ToJson};
use rustc_serialize::json::DecoderError::MissingFieldError;
use rustc_serialize::Decodable;
use rustc_serialize::Decoder;

use slack::Message;

use adapter::ChatAdapter;
use message::AdapterMsg;
use message::IncomingMessage;
use message::OutgoingMessage;

/// SlackAdapter sends and receives messages from the Slack chat service. Until actualy
/// configuration is added, the slack token should be placed in the environment variable
/// `SLACK_BOT_TOKEN`
pub struct SlackAdapter {
    token: String
}

impl SlackAdapter {
    pub fn new() -> SlackAdapter {
        SlackAdapter {
            token: match env::var("SLACK_BOT_TOKEN") {
                Ok(t) => t,
                Err(_) => panic!("Failed to get SLACK_BOT_TOKEN from env")
            }
        }
    }
}

struct MyHandler {
  count: i64,
  tx_bot: Sender<IncomingMessage>,
  tx_adapter: Sender<AdapterMsg>
}

/// Data for a Event::Message
#[allow(dead_code)]
struct MessageData {
    text: String,
    channel: String,
    user: String,
    ts: String,
    team: String
}

macro_rules! str_accessor {
    ($s:ident) => {
        pub fn $s(&self) -> &str {
            self.$s.as_ref()
        }
    }
}

#[allow(dead_code)]
impl MessageData {
    str_accessor!(text);
    str_accessor!(channel);
    str_accessor!(user);
    str_accessor!(ts);
    str_accessor!(team);
}

/// Incoming slack messages on the websocket api
enum Event {
    /// A message was sent to a channel
    Message(Msg),

    /// Some other type of message arrived.
    /// [The list](https://api.slack.com/rtm) is quite extensive, and only the
    /// messages the adapter is concerned with are enumerated here.
    Other(json::Object)
}

/// Event::Message sub types. Message events are the only event capable of having a sub type.
enum Msg {
    /// A regular text message from a user
    Plain(MessageData),

    /// Subtype not explicitly handled. For a complete enumeration, check
    /// https://api.slack.com/events/message
    Other(json::Object)
}

#[derive(Debug)]
struct MessageDecodeError(&'static str);

impl From<json::BuilderError> for MessageDecodeError {
    fn from(err: json::BuilderError) -> MessageDecodeError {
        MessageDecodeError("json::BuilderError")
    }
}

/// Extract a &str ($key) from a json::Object ($obj)
/// Returns an Err(MessageDecodeError) when it fails
macro_rules! get_json_string {
    ($obj:ident, $key:expr) => {
        {
            let json_str: &Json = match $obj.get($key) {
                Some(json_str) => json_str,
                None => return Err(MessageDecodeError("missing field"))
            };

            match json_str.as_string() {
                Some(slice) => slice.to_owned(),
                None => return Err(MessageDecodeError("not a string"))
            }
        }
    }
}

fn decode_msg_json(obj: json::Object, msg_type: &str) -> Result<Event, MessageDecodeError> {
    // Messages with a `subtype` are not plain text messages.. Not interested in them for now.
    if obj.contains_key("subtype") {
        return Ok(Event::Message(Msg::Other(obj)))
    }

    // It's a plain message.
    Ok(Event::Message(Msg::Plain(MessageData {
        text: get_json_string!(obj, "text"),
        channel: get_json_string!(obj, "channel"),
        user: get_json_string!(obj, "user"),
        ts: get_json_string!(obj, "ts"),
        team: get_json_string!(obj, "team"),
    })))
}

fn decode_json(raw: &str) -> Result<Event, MessageDecodeError> {
    let json = try!(Json::from_str(raw));
    let obj = match json {
        Json::Object(obj) => obj,
        _ => return Err(MessageDecodeError("string is not json object"))
    };

    // If the message does not have a `type`, it is a confirmation message
    if !obj.contains_key("type") {
        return Ok(Event::Other(obj))
    }

    let msg = match obj.get("type").expect("obj has key type but failed to get").as_string() {
        Some(s) => s.to_owned(),
        None => return Err(MessageDecodeError("type field is not a string"))
    };

    match msg.as_ref() {
        // If it's a message type, try and return an Event::Message
        "message" => decode_msg_json(obj, msg.as_ref()),
        // Some other type we don't explicitly handle
        _ => Ok(Event::Other(obj)),
    }
}

/// Convert a JSON string to a Event
///
/// This methods provides additional error handling around json::decode for certain errors
/// that cannot be handled in the Decodable implementation. Specifically, MissingFieldError
/// where the field is "type" are actually valid messages despite missing the "type" field.
fn string_to_slack_msg(raw: &str) -> Result<Event, MessageDecodeError> {
    // Some messages arriving from the slack client don't have a type. So far I've only
    // witnessed confirmation messages arriving in this fashion. Since they go through the same
    // pipeline as content messages, the decoder should be able to handle them.
    decode_json(raw)
}

#[derive(Debug)]
struct OutgoingEvent {
    id: i64,
    channel: String,
    msg_type: String,
    text: String
}

impl OutgoingEvent {
    fn new(id: i64, m: OutgoingMessage) -> OutgoingEvent {
        OutgoingEvent {
            id: id,
            channel: m.get_incoming().channel().expect("missing channel").to_owned(),
            msg_type: "message".to_owned(),
            text: m.as_ref().to_owned() // TODO move instead of copy
        }
    }
}

impl ToJson for OutgoingEvent {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        d.insert("id".to_string(), self.id.to_json());
        d.insert("channel".to_string(), self.channel.to_json());
        d.insert("type".to_string(), self.msg_type.to_json());
        d.insert("text".to_string(), self.text.to_json());
        Json::Object(d)
    }
}

#[allow(unused_variables)]
impl slack::MessageHandler for MyHandler {
    fn on_receive(&mut self, cli: &mut slack::RtmClient, raw: &str) {
        println!("Received[{}]: {}", self.count, raw.to_string());
        self.count = self.count + 1;

        match string_to_slack_msg(raw) {
            Ok(slack_msg) => {
                match slack_msg {
                    Event::Message(Msg::Plain(msg)) => {
                        let incoming = IncomingMessage::new("SlackAdapter".to_owned(), None,
                            Some(msg.channel().to_owned()), Some(msg.user().to_owned()),
                            msg.text().to_owned(), self.tx_adapter.clone());

                        self.tx_bot.send(incoming).ok().expect("Bot unable to process messages");
                    },
                    _ => ()
                }
            },
            Err(e) => {
                println!("error decoding slack message: {:?}", e);
                println!("please consider reporting this to jwilm/chatbot as it is probably a bug");
            }
        }
    }

    fn on_ping(&mut self, cli: &mut slack::RtmClient) { }

    fn on_close(&mut self, cli: &mut slack::RtmClient) { }

    fn on_connect(&mut self, cli: &mut slack::RtmClient) { }
}

impl ChatAdapter for SlackAdapter {
    /// SlackAdapter name
    fn get_name(&self) -> &str {
        "SlackAdapter"
    }

    fn process_events(&self) -> Receiver<IncomingMessage> {
        println!("SlackAdapter: process_events");
        let (tx_bot, rx_bot) = channel();
        let (tx_adapter, rx_adapter) = channel();

        let uid = AtomicIsize::new(0);

        let mut cli = slack::RtmClient::new();
        let (client, slack_rx) = cli.login(self.token.as_ref()).unwrap();
        let slack_tx = cli.get_outs().unwrap();

        thread::Builder::new().name("Chatbot Slack Receiver".to_owned()).spawn(move || {
            let mut handler = MyHandler{count: 0, tx_bot: tx_bot, tx_adapter: tx_adapter};
            cli.run::<MyHandler>(&mut handler, client, slack_rx).unwrap();
        }).ok().expect("failed to create thread for slack receiver");

        thread::Builder::new().name("Chatbot Slack Sender".to_owned()).spawn(move || {
            loop {
                match rx_adapter.recv() {
                    Ok(msg) => {
                        match msg {
                            AdapterMsg::Outgoing(m) => {
                                let id = uid.fetch_add(1, Ordering::SeqCst) as i64;
                                let out = OutgoingEvent::new(id, m);
                                slack_tx.send(Message::Text(out.to_json().to_string())).unwrap();
                            }
                            _ => unreachable!("No other messages being sent yet")
                        }
                    },
                    Err(e) => {
                        println!("error receiving outgoing messages: {}", e);
                        break
                    }
                }
            }
        }).ok().expect("failed to create thread for slack sender");

        rx_bot
    }
}

#[cfg(test)]
mod tests {
    use adapter::slack::Event;
    use adapter::slack::Msg;
    use adapter::slack::string_to_slack_msg;

    #[test]
    fn decode_message() {
        let raw = "{\"type\":\"message\", \
                    \"channel\":\"D04UYUAMW\", \
                    \"user\":\"U02ALMR84\", \
                    \"text\":\"ping\", \
                    \"ts\":\"1432563914.000007\", \
                    \"team\":\"T02ALMR82\"}";

        let slack_msg = string_to_slack_msg(raw).unwrap();

        match slack_msg {
            Event::Message(Msg::Plain(data)) => {
                assert_eq!(data.text(), "ping");
                assert_eq!(data.channel(), "D04UYUAMW");
                assert_eq!(data.user(), "U02ALMR84");
                assert_eq!(data.ts(), "1432563914.000007");
                assert_eq!(data.team(), "T02ALMR82");
            },
            _ => panic!("Expected Event::Message")
        }
    }
    #[test]
    fn decode_confirmation_message() {
        let raw = r#"{"ok":true,"reply_to":0,"ts":"1432566639.000014","text":"pong"}"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Other(_) => (),
            _ => panic!("expected Event::Other")
        }
    }

    #[test]
    fn decode_unexpected_type() {
        let raw = r#"{"type":"not_a_slack_msg_type"}"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Other(_) => (),
            _ => panic!("expected Event::Other")
        }
    }

    #[test]
    fn decode_message_changed() {
        let raw = r#"{
            "type":"message",
            "message":{
                "type":"message",
                "user":"U02ALMR84",
                "text":"arst",
                "edited":{
                    "user":"U02ALMR84",
                    "ts":"1432695814.000000"
                },
                "ts":"1432695812.000056"
            },
            "subtype":"message_changed",
            "hidden":true,
            "channel":"D04UYUAMW",
            "event_ts":"1432695814.616510",
            "ts":"1432695814.000057"
        }"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Message(Msg::Other(_)) => return,
            _ => panic!("expected Msg::Other")
        }
    }

    #[test]
    fn decode_me_message() {
        let raw = r#"{
            "text":"is a potato",
            "type":"message",
            "subtype":"me_message",
            "user":"U02ALMR84",
            "channel":"D04UYUAMW",
            "ts":"1432695826.000060"
        }"#;
    }

    #[test]
    fn decode_message_deleted() {
        let raw = r#"{
            "type":"message",
            "deleted_ts":"1432695826.000060",
            "subtype":"message_deleted",
            "hidden":true,
            "channel":"D04UYUAMW",
            "event_ts":"1432695848.617155",
            "ts":"1432695848.000061"
        }"#;
    }

}
<commit_msg>slack: finish implementing message subtype tests<commit_after>extern crate slack;

use std::collections::BTreeMap;
use std::env;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;
use std::sync::mpsc::channel;
use std::thread;

use rustc_serialize::json::{self, Json, ToJson};

use slack::Message;

use adapter::ChatAdapter;
use message::AdapterMsg;
use message::IncomingMessage;
use message::OutgoingMessage;

/// SlackAdapter sends and receives messages from the Slack chat service. Until actualy
/// configuration is added, the slack token should be placed in the environment variable
/// `SLACK_BOT_TOKEN`
pub struct SlackAdapter {
    token: String
}

impl SlackAdapter {
    pub fn new() -> SlackAdapter {
        SlackAdapter {
            token: match env::var("SLACK_BOT_TOKEN") {
                Ok(t) => t,
                Err(_) => panic!("Failed to get SLACK_BOT_TOKEN from env")
            }
        }
    }
}

struct MyHandler {
  count: i64,
  tx_bot: Sender<IncomingMessage>,
  tx_adapter: Sender<AdapterMsg>
}

/// Data for a Event::Message
#[allow(dead_code)]
struct MessageData {
    text: String,
    channel: String,
    user: String,
    ts: String,
    team: String
}

macro_rules! str_accessor {
    ($s:ident) => {
        pub fn $s(&self) -> &str {
            self.$s.as_ref()
        }
    }
}

#[allow(dead_code)]
impl MessageData {
    str_accessor!(text);
    str_accessor!(channel);
    str_accessor!(user);
    str_accessor!(ts);
    str_accessor!(team);
}

/// Incoming slack messages on the websocket api
enum Event {
    /// A message was sent to a channel
    Message(Msg),

    /// Some other type of message arrived.
    /// [The list](https://api.slack.com/rtm) is quite extensive, and only the
    /// messages the adapter is concerned with are enumerated here.
    Other(json::Object)
}

/// Event::Message sub types. Message events are the only event capable of having a sub type.
enum Msg {
    /// A regular text message from a user
    Plain(MessageData),

    /// Subtype not explicitly handled. For a complete enumeration, check
    /// https://api.slack.com/events/message
    Other(json::Object)
}

#[derive(Debug)]
struct MessageDecodeError(&'static str);

impl From<json::BuilderError> for MessageDecodeError {
    fn from(_: json::BuilderError) -> MessageDecodeError {
        MessageDecodeError("json::BuilderError")
    }
}

/// Extract a &str ($key) from a json::Object ($obj)
/// Returns an Err(MessageDecodeError) when it fails
macro_rules! get_json_string {
    ($obj:ident, $key:expr) => {
        {
            let json_str: &Json = match $obj.get($key) {
                Some(json_str) => json_str,
                None => return Err(MessageDecodeError("missing field"))
            };

            match json_str.as_string() {
                Some(slice) => slice.to_owned(),
                None => return Err(MessageDecodeError("not a string"))
            }
        }
    }
}

fn decode_msg_json(obj: json::Object) -> Result<Event, MessageDecodeError> {
    // Messages with a `subtype` are not plain text messages.. Not interested in them for now.
    if obj.contains_key("subtype") {
        return Ok(Event::Message(Msg::Other(obj)))
    }

    // It's a plain message.
    Ok(Event::Message(Msg::Plain(MessageData {
        text: get_json_string!(obj, "text"),
        channel: get_json_string!(obj, "channel"),
        user: get_json_string!(obj, "user"),
        ts: get_json_string!(obj, "ts"),
        team: get_json_string!(obj, "team"),
    })))
}

fn decode_json(raw: &str) -> Result<Event, MessageDecodeError> {
    let json = try!(Json::from_str(raw));
    let obj = match json {
        Json::Object(obj) => obj,
        _ => return Err(MessageDecodeError("string is not json object"))
    };

    // If the message does not have a `type`, it is a confirmation message
    if !obj.contains_key("type") {
        return Ok(Event::Other(obj))
    }

    let msg = match obj.get("type").expect("obj has key type but failed to get").as_string() {
        Some(s) => s.to_owned(),
        None => return Err(MessageDecodeError("type field is not a string"))
    };

    match msg.as_ref() {
        // If it's a message type, try and return an Event::Message
        "message" => decode_msg_json(obj),
        // Some other type we don't explicitly handle
        _ => Ok(Event::Other(obj)),
    }
}

/// Convert a JSON string to a Event
///
/// This methods provides additional error handling around json::decode for certain errors
/// that cannot be handled in the Decodable implementation. Specifically, MissingFieldError
/// where the field is "type" are actually valid messages despite missing the "type" field.
fn string_to_slack_msg(raw: &str) -> Result<Event, MessageDecodeError> {
    // Some messages arriving from the slack client don't have a type. So far I've only
    // witnessed confirmation messages arriving in this fashion. Since they go through the same
    // pipeline as content messages, the decoder should be able to handle them.
    decode_json(raw)
}

#[derive(Debug)]
struct OutgoingEvent {
    id: i64,
    channel: String,
    msg_type: String,
    text: String
}

impl OutgoingEvent {
    fn new(id: i64, m: OutgoingMessage) -> OutgoingEvent {
        OutgoingEvent {
            id: id,
            channel: m.get_incoming().channel().expect("missing channel").to_owned(),
            msg_type: "message".to_owned(),
            text: m.as_ref().to_owned() // TODO move instead of copy
        }
    }
}

impl ToJson for OutgoingEvent {
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        d.insert("id".to_string(), self.id.to_json());
        d.insert("channel".to_string(), self.channel.to_json());
        d.insert("type".to_string(), self.msg_type.to_json());
        d.insert("text".to_string(), self.text.to_json());
        Json::Object(d)
    }
}

#[allow(unused_variables)]
impl slack::MessageHandler for MyHandler {
    fn on_receive(&mut self, cli: &mut slack::RtmClient, raw: &str) {
        println!("Received[{}]: {}", self.count, raw.to_string());
        self.count = self.count + 1;

        match string_to_slack_msg(raw) {
            Ok(slack_msg) => {
                match slack_msg {
                    Event::Message(Msg::Plain(msg)) => {
                        let incoming = IncomingMessage::new("SlackAdapter".to_owned(), None,
                            Some(msg.channel().to_owned()), Some(msg.user().to_owned()),
                            msg.text().to_owned(), self.tx_adapter.clone());

                        self.tx_bot.send(incoming).ok().expect("Bot unable to process messages");
                    },
                    _ => ()
                }
            },
            Err(e) => {
                println!("error decoding slack message: {:?}", e);
                println!("please consider reporting this to jwilm/chatbot as it is probably a bug");
            }
        }
    }

    fn on_ping(&mut self, cli: &mut slack::RtmClient) { }

    fn on_close(&mut self, cli: &mut slack::RtmClient) { }

    fn on_connect(&mut self, cli: &mut slack::RtmClient) { }
}

impl ChatAdapter for SlackAdapter {
    /// SlackAdapter name
    fn get_name(&self) -> &str {
        "SlackAdapter"
    }

    fn process_events(&self) -> Receiver<IncomingMessage> {
        println!("SlackAdapter: process_events");
        let (tx_bot, rx_bot) = channel();
        let (tx_adapter, rx_adapter) = channel();

        let uid = AtomicIsize::new(0);

        let mut cli = slack::RtmClient::new();
        let (client, slack_rx) = cli.login(self.token.as_ref()).unwrap();
        let slack_tx = cli.get_outs().unwrap();

        thread::Builder::new().name("Chatbot Slack Receiver".to_owned()).spawn(move || {
            let mut handler = MyHandler{count: 0, tx_bot: tx_bot, tx_adapter: tx_adapter};
            cli.run::<MyHandler>(&mut handler, client, slack_rx).unwrap();
        }).ok().expect("failed to create thread for slack receiver");

        thread::Builder::new().name("Chatbot Slack Sender".to_owned()).spawn(move || {
            loop {
                match rx_adapter.recv() {
                    Ok(msg) => {
                        match msg {
                            AdapterMsg::Outgoing(m) => {
                                let id = uid.fetch_add(1, Ordering::SeqCst) as i64;
                                let out = OutgoingEvent::new(id, m);
                                slack_tx.send(Message::Text(out.to_json().to_string())).unwrap();
                            }
                            _ => unreachable!("No other messages being sent yet")
                        }
                    },
                    Err(e) => {
                        println!("error receiving outgoing messages: {}", e);
                        break
                    }
                }
            }
        }).ok().expect("failed to create thread for slack sender");

        rx_bot
    }
}

#[cfg(test)]
mod tests {
    use adapter::slack::Event;
    use adapter::slack::Msg;
    use adapter::slack::string_to_slack_msg;

    #[test]
    fn decode_message() {
        let raw = "{\"type\":\"message\", \
                    \"channel\":\"D04UYUAMW\", \
                    \"user\":\"U02ALMR84\", \
                    \"text\":\"ping\", \
                    \"ts\":\"1432563914.000007\", \
                    \"team\":\"T02ALMR82\"}";

        let slack_msg = string_to_slack_msg(raw).unwrap();

        match slack_msg {
            Event::Message(Msg::Plain(data)) => {
                assert_eq!(data.text(), "ping");
                assert_eq!(data.channel(), "D04UYUAMW");
                assert_eq!(data.user(), "U02ALMR84");
                assert_eq!(data.ts(), "1432563914.000007");
                assert_eq!(data.team(), "T02ALMR82");
            },
            _ => panic!("Expected Event::Message")
        }
    }
    #[test]
    fn decode_confirmation_message() {
        let raw = r#"{"ok":true,"reply_to":0,"ts":"1432566639.000014","text":"pong"}"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Other(_) => (),
            _ => panic!("expected Event::Other")
        }
    }

    #[test]
    fn decode_unexpected_type() {
        let raw = r#"{"type":"not_a_slack_msg_type"}"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Other(_) => (),
            _ => panic!("expected Event::Other")
        }
    }

    #[test]
    fn decode_message_changed() {
        let raw = r#"{
            "type":"message",
            "message":{
                "type":"message",
                "user":"U02ALMR84",
                "text":"arst",
                "edited":{
                    "user":"U02ALMR84",
                    "ts":"1432695814.000000"
                },
                "ts":"1432695812.000056"
            },
            "subtype":"message_changed",
            "hidden":true,
            "channel":"D04UYUAMW",
            "event_ts":"1432695814.616510",
            "ts":"1432695814.000057"
        }"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Message(Msg::Other(_)) => return,
            _ => panic!("expected Msg::Other")
        }
    }

    #[test]
    fn decode_me_message() {
        let raw = r#"{
            "text":"is a potato",
            "type":"message",
            "subtype":"me_message",
            "user":"U02ALMR84",
            "channel":"D04UYUAMW",
            "ts":"1432695826.000060"
        }"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Message(Msg::Other(_)) => return,
            _ => panic!("expected Msg::Other")
        }
    }

    #[test]
    fn decode_message_deleted() {
        let raw = r#"{
            "type":"message",
            "deleted_ts":"1432695826.000060",
            "subtype":"message_deleted",
            "hidden":true,
            "channel":"D04UYUAMW",
            "event_ts":"1432695848.617155",
            "ts":"1432695848.000061"
        }"#;

        match string_to_slack_msg(raw).unwrap() {
            Event::Message(Msg::Other(_)) => return,
            _ => panic!("expected Msg::Other")
        }
    }

}
<|endoftext|>


<commit_before>use crate::api::internal::InterConfig;
use crate::config::EncoderConfig;
use crate::context::{BlockOffset, FrameBlocks, TileBlockOffset};
use crate::cpu_features::CpuFeatureLevel;
use crate::dist::get_satd;
use crate::encoder::{
  FrameInvariants, FrameState, Sequence, IMPORTANCE_BLOCK_SIZE,
};
use crate::frame::{AsRegion, PlaneOffset};
use crate::me::estimate_tile_motion;
use crate::partition::{get_intra_edges, BlockSize};
use crate::predict::{IntraParam, PredictionMode};
use crate::rayon::iter::*;
use crate::tiling::{Area, TileRect};
use crate::transform::TxSize;
use crate::{Frame, Pixel};
use histogram::Histogram;
use rust_hawktracer::*;
use std::sync::Arc;
use v_frame::pixel::CastFromPrimitive;

pub(crate) const IMP_BLOCK_MV_UNITS_PER_PIXEL: i64 = 8;
pub(crate) const IMP_BLOCK_SIZE_IN_MV_UNITS: i64 =
  IMPORTANCE_BLOCK_SIZE as i64 * IMP_BLOCK_MV_UNITS_PER_PIXEL;
pub(crate) const IMP_BLOCK_AREA_IN_MV_UNITS: i64 =
  IMP_BLOCK_SIZE_IN_MV_UNITS * IMP_BLOCK_SIZE_IN_MV_UNITS;

#[hawktracer(estimate_intra_costs)]
pub(crate) fn estimate_intra_costs<T: Pixel>(
  frame: &Frame<T>, bit_depth: usize, cpu_feature_level: CpuFeatureLevel,
) -> Box<[u32]> {
  let plane = &frame.planes[0];
  let mut plane_after_prediction = frame.planes[0].clone();

  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  let tx_size = bsize.tx_size();

  let h_in_imp_b = plane.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut intra_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);

  for y in 0..h_in_imp_b {
    for x in 0..w_in_imp_b {
      let plane_org = plane.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      // TODO: other intra prediction modes.
      let edge_buf = get_intra_edges(
        &plane.as_region(),
        TileBlockOffset(BlockOffset { x, y }),
        0,
        0,
        bsize,
        PlaneOffset {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        },
        TxSize::TX_8X8,
        bit_depth,
        Some(PredictionMode::DC_PRED),
        false,
        IntraParam::None,
      );

      let mut plane_after_prediction_region = plane_after_prediction
        .region_mut(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      PredictionMode::DC_PRED.predict_intra(
        TileRect {
          x: x * IMPORTANCE_BLOCK_SIZE,
          y: y * IMPORTANCE_BLOCK_SIZE,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        },
        &mut plane_after_prediction_region,
        tx_size,
        bit_depth,
        &[], // Not used by DC_PRED
        IntraParam::None,
        None, // Not used by DC_PRED
        &edge_buf,
        cpu_feature_level,
      );

      let plane_after_prediction_region =
        plane_after_prediction.region(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      let intra_cost = get_satd(
        &plane_org,
        &plane_after_prediction_region,
        bsize,
        bit_depth,
        cpu_feature_level,
      );

      intra_costs.push(intra_cost);
    }
  }

  intra_costs.into_boxed_slice()
}

#[hawktracer(estimate_inter_costs_histogram)]
pub(crate) fn estimate_inter_costs_histogram_blocks<T: Pixel>(
  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>,
) -> Box<[u32]> {
  let plane_org = &frame.planes[0];
  let plane_ref = &ref_frame.planes[0];
  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut inter_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);

  (0..h_in_imp_b).for_each(|y| {
    (0..w_in_imp_b).for_each(|x| {
      // Coordinates of the top-left corner of the reference block, in MV
      // units.
      let region_org = plane_org.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let region_ref = plane_ref.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let mut histogram_org =
        Histogram::configure().max_value(256).build().unwrap();
      let iter_org = region_org.rows_iter();
      for row in iter_org {
        for pixel in row {
          let cur = i16::cast_from(*pixel);
          histogram_org.increment(cur as u64).unwrap();
        }
      }

      let mut histogram_ref =
        Histogram::configure().max_value(256).build().unwrap();
      let iter_ref = region_ref.rows_iter();
      for row in iter_ref {
        for pixel in row {
          let cur = i16::cast_from(*pixel);
          histogram_ref.increment(cur as u64).unwrap();
        }
      }

      let mean = (histogram_org.mean().unwrap() as i32
        - histogram_ref.mean().unwrap() as i32)
        .abs();

      inter_costs.push(mean as u32);
    });
  });
  inter_costs.into_boxed_slice()
}

#[hawktracer(estimate_inter_costs)]
pub(crate) fn estimate_inter_costs<T: Pixel>(
  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>, bit_depth: usize,
  mut config: EncoderConfig, sequence: Arc<Sequence>,
) -> Box<[u32]> {
  config.low_latency = true;
  config.speed_settings.multiref = false;
  let inter_cfg = InterConfig::new(&config);
  let last_fi = FrameInvariants::new_key_frame(Arc::new(config), sequence, 0);
  let mut fi =
    FrameInvariants::new_inter_frame(&last_fi, &inter_cfg, 0, 1, 2, false);

  // Compute the motion vectors.
  let mut fs = FrameState::new_with_frame(&fi, frame.clone());
  compute_motion_vectors(&mut fi, &mut fs, &inter_cfg);

  // Estimate inter costs
  let plane_org = &frame.planes[0];
  let plane_ref = &ref_frame.planes[0];
  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut inter_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);
  let stats = &fs.frame_me_stats[0];
  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  (0..h_in_imp_b).for_each(|y| {
    (0..w_in_imp_b).for_each(|x| {
      let mv = stats[y * 2][x * 2].mv;

      // Coordinates of the top-left corner of the reference block, in MV
      // units.
      let reference_x = x as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.col as i64;
      let reference_y = y as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.row as i64;

      let region_org = plane_org.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let region_ref = plane_ref.region(Area::Rect {
        x: reference_x as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        y: reference_y as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      inter_costs.push(get_satd(
        &region_org,
        &region_ref,
        bsize,
        bit_depth,
        fi.cpu_feature_level,
      ));
    });
  });
  inter_costs.into_boxed_slice()
}

#[hawktracer(compute_motion_vectors)]
pub(crate) fn compute_motion_vectors<T: Pixel>(
  fi: &mut FrameInvariants<T>, fs: &mut FrameState<T>, inter_cfg: &InterConfig,
) {
  let mut blocks = FrameBlocks::new(fi.w_in_b, fi.h_in_b);
  fi.sequence
    .tiling
    .tile_iter_mut(fs, &mut blocks)
    .collect::<Vec<_>>()
    .into_par_iter()
    .for_each(|mut ctx| {
      let ts = &mut ctx.ts;
      estimate_tile_motion(fi, ts, inter_cfg);
    });
}
<commit_msg>Fix the histogram max_value setting<commit_after>use crate::api::internal::InterConfig;
use crate::config::EncoderConfig;
use crate::context::{BlockOffset, FrameBlocks, TileBlockOffset};
use crate::cpu_features::CpuFeatureLevel;
use crate::dist::get_satd;
use crate::encoder::{
  FrameInvariants, FrameState, Sequence, IMPORTANCE_BLOCK_SIZE,
};
use crate::frame::{AsRegion, PlaneOffset};
use crate::me::estimate_tile_motion;
use crate::partition::{get_intra_edges, BlockSize};
use crate::predict::{IntraParam, PredictionMode};
use crate::rayon::iter::*;
use crate::tiling::{Area, TileRect};
use crate::transform::TxSize;
use crate::{Frame, Pixel};
use histogram::Histogram;
use rust_hawktracer::*;
use std::sync::Arc;
use v_frame::pixel::CastFromPrimitive;

pub(crate) const IMP_BLOCK_MV_UNITS_PER_PIXEL: i64 = 8;
pub(crate) const IMP_BLOCK_SIZE_IN_MV_UNITS: i64 =
  IMPORTANCE_BLOCK_SIZE as i64 * IMP_BLOCK_MV_UNITS_PER_PIXEL;
pub(crate) const IMP_BLOCK_AREA_IN_MV_UNITS: i64 =
  IMP_BLOCK_SIZE_IN_MV_UNITS * IMP_BLOCK_SIZE_IN_MV_UNITS;

#[hawktracer(estimate_intra_costs)]
pub(crate) fn estimate_intra_costs<T: Pixel>(
  frame: &Frame<T>, bit_depth: usize, cpu_feature_level: CpuFeatureLevel,
) -> Box<[u32]> {
  let plane = &frame.planes[0];
  let mut plane_after_prediction = frame.planes[0].clone();

  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  let tx_size = bsize.tx_size();

  let h_in_imp_b = plane.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut intra_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);

  for y in 0..h_in_imp_b {
    for x in 0..w_in_imp_b {
      let plane_org = plane.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      // TODO: other intra prediction modes.
      let edge_buf = get_intra_edges(
        &plane.as_region(),
        TileBlockOffset(BlockOffset { x, y }),
        0,
        0,
        bsize,
        PlaneOffset {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        },
        TxSize::TX_8X8,
        bit_depth,
        Some(PredictionMode::DC_PRED),
        false,
        IntraParam::None,
      );

      let mut plane_after_prediction_region = plane_after_prediction
        .region_mut(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      PredictionMode::DC_PRED.predict_intra(
        TileRect {
          x: x * IMPORTANCE_BLOCK_SIZE,
          y: y * IMPORTANCE_BLOCK_SIZE,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        },
        &mut plane_after_prediction_region,
        tx_size,
        bit_depth,
        &[], // Not used by DC_PRED
        IntraParam::None,
        None, // Not used by DC_PRED
        &edge_buf,
        cpu_feature_level,
      );

      let plane_after_prediction_region =
        plane_after_prediction.region(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      let intra_cost = get_satd(
        &plane_org,
        &plane_after_prediction_region,
        bsize,
        bit_depth,
        cpu_feature_level,
      );

      intra_costs.push(intra_cost);
    }
  }

  intra_costs.into_boxed_slice()
}

#[hawktracer(estimate_inter_costs_histogram)]
pub(crate) fn estimate_inter_costs_histogram_blocks<T: Pixel>(
  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>,
) -> Box<[u32]> {
  let plane_org = &frame.planes[0];
  let plane_ref = &ref_frame.planes[0];
  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut inter_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);

  (0..h_in_imp_b).for_each(|y| {
    (0..w_in_imp_b).for_each(|x| {
      // Coordinates of the top-left corner of the reference block, in MV
      // units.
      let region_org = plane_org.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let region_ref = plane_ref.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let max_value = 1u64 << 12; // At most 12 bit per pixel

      let mut histogram_org =
        Histogram::configure().max_value(max_value).build().unwrap();
      let iter_org = region_org.rows_iter();
      for row in iter_org {
        for pixel in row {
          let cur = i16::cast_from(*pixel);
          histogram_org.increment(cur as u64).unwrap();
        }
      }

      let mut histogram_ref =
        Histogram::configure().max_value(max_value).build().unwrap();
      let iter_ref = region_ref.rows_iter();
      for row in iter_ref {
        for pixel in row {
          let cur = i16::cast_from(*pixel);
          histogram_ref.increment(cur as u64).unwrap();
        }
      }

      let mean = (histogram_org.mean().unwrap() as i32
        - histogram_ref.mean().unwrap() as i32)
        .abs();

      inter_costs.push(mean as u32);
    });
  });
  inter_costs.into_boxed_slice()
}

#[hawktracer(estimate_inter_costs)]
pub(crate) fn estimate_inter_costs<T: Pixel>(
  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>, bit_depth: usize,
  mut config: EncoderConfig, sequence: Arc<Sequence>,
) -> Box<[u32]> {
  config.low_latency = true;
  config.speed_settings.multiref = false;
  let inter_cfg = InterConfig::new(&config);
  let last_fi = FrameInvariants::new_key_frame(Arc::new(config), sequence, 0);
  let mut fi =
    FrameInvariants::new_inter_frame(&last_fi, &inter_cfg, 0, 1, 2, false);

  // Compute the motion vectors.
  let mut fs = FrameState::new_with_frame(&fi, frame.clone());
  compute_motion_vectors(&mut fi, &mut fs, &inter_cfg);

  // Estimate inter costs
  let plane_org = &frame.planes[0];
  let plane_ref = &ref_frame.planes[0];
  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut inter_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);
  let stats = &fs.frame_me_stats[0];
  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  (0..h_in_imp_b).for_each(|y| {
    (0..w_in_imp_b).for_each(|x| {
      let mv = stats[y * 2][x * 2].mv;

      // Coordinates of the top-left corner of the reference block, in MV
      // units.
      let reference_x = x as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.col as i64;
      let reference_y = y as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.row as i64;

      let region_org = plane_org.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let region_ref = plane_ref.region(Area::Rect {
        x: reference_x as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        y: reference_y as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      inter_costs.push(get_satd(
        &region_org,
        &region_ref,
        bsize,
        bit_depth,
        fi.cpu_feature_level,
      ));
    });
  });
  inter_costs.into_boxed_slice()
}

#[hawktracer(compute_motion_vectors)]
pub(crate) fn compute_motion_vectors<T: Pixel>(
  fi: &mut FrameInvariants<T>, fs: &mut FrameState<T>, inter_cfg: &InterConfig,
) {
  let mut blocks = FrameBlocks::new(fi.w_in_b, fi.h_in_b);
  fi.sequence
    .tiling
    .tile_iter_mut(fs, &mut blocks)
    .collect::<Vec<_>>()
    .into_par_iter()
    .for_each(|mut ctx| {
      let ts = &mut ctx.ts;
      estimate_tile_motion(fi, ts, inter_cfg);
    });
}
<|endoftext|>


<commit_before>//! # Jaeger JSON file Exporter
//!

use async_trait::async_trait;
use futures::{future::BoxFuture, FutureExt};
use opentelemetry::sdk::export::trace::{ExportResult, SpanData, SpanExporter};
use opentelemetry::sdk::trace::{TraceRuntime, Tracer};
use opentelemetry::trace::TraceError;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::SystemTime;

/// An exporter for jaeger comptible json files containing trace data
#[derive(Debug)]
pub struct JaegerJsonExporter<R> {
    out_path: PathBuf,
    file_prefix: String,
    service_name: String,
    runtime: R,
}

impl<R: JaegerJsonRuntime> JaegerJsonExporter<R> {
    /// Configure a new jaeger-json exporter
    ///
    /// * `out_path` refers to an directory where span data are written. If it does not exist, it is created by the exporter
    /// * `file_prefix` refers to a prefix prependend to each span file
    /// * `service_name` is used to identify the corresponding service in jaeger
    /// * `runtime` specifies the used async runtime to write the trace data
    pub fn new(out_path: PathBuf, file_prefix: String, service_name: String, runtime: R) -> Self {
        Self {
            out_path,
            file_prefix,
            service_name,
            runtime,
        }
    }

    /// Install the exporter using the internal provided runtime
    pub fn install_batch(self) -> Tracer {
        use opentelemetry::trace::TracerProvider;

        let runtime = self.runtime.clone();
        let provider_builder =
            opentelemetry::sdk::trace::TracerProvider::builder().with_batch_exporter(self, runtime);

        let provider = provider_builder.build();

        let tracer =
            provider.versioned_tracer("opentelemetry", Some(env!("CARGO_PKG_VERSION")), None);
        let _ = opentelemetry::global::set_tracer_provider(provider);

        tracer
    }
}

impl<R: JaegerJsonRuntime> SpanExporter for JaegerJsonExporter<R> {
    fn export(&mut self, batch: Vec<SpanData>) -> BoxFuture<'static, ExportResult> {
        let mut trace_map = HashMap::new();

        for span in batch {
            let ctx = &span.span_context;
            trace_map
                .entry(ctx.trace_id())
                .or_insert_with(Vec::new)
                .push(span_data_to_jaeger_json(span));
        }

        let data = trace_map
            .into_iter()
            .map(|(trace_id, spans)| {
                serde_json::json!({
                    "traceID": trace_id.to_string(),
                    "spans": spans,
                    "processes": {
                        "p1": {
                            "serviceName": self.service_name,
                            "tags": []
                        }
                    }
                })
            })
            .collect::<Vec<_>>();

        let json = serde_json::json!({
            "data": data,
        });

        let runtime = self.runtime.clone();
        let out_path = self.out_path.clone();
        let file_prefix = self.file_prefix.clone();

        async move {
            runtime.create_dir(&out_path).await?;

            let file_name = out_path.join(format!(
                "{}-{}.json",
                file_prefix,
                SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .expect("This does not fail")
                    .as_secs()
            ));
            runtime
                .write_to_file(
                    &file_name,
                    &serde_json::to_vec(&json).expect("This is a valid json value"),
                )
                .await?;

            Ok(())
        }
        .boxed()
    }
}

fn span_data_to_jaeger_json(
    span: opentelemetry::sdk::export::trace::SpanData,
) -> serde_json::Value {
    let events = span
        .events
        .iter()
        .map(|e| {
            let mut fields = e
                .attributes
                .iter()
                .map(|a| {
                    let (tpe, value) = opentelemetry_value_to_json(&a.value);
                    serde_json::json!({
                        "key": a.key.as_str(),
                        "type": tpe,
                        "value": value,
                    })
                })
                .collect::<Vec<_>>();
            fields.push(serde_json::json!({
                "key": "event",
                "type": "string",
                "value": e.name,
            }));

            serde_json::json!({
                "timestamp": e.timestamp.duration_since(SystemTime::UNIX_EPOCH).expect("This does not fail").as_micros() as i64,
                "fields": fields,
            })
        })
        .collect::<Vec<_>>();
    let tags = span
        .attributes
        .iter()
        .map(|(key, value)| {
            let (tpe, value) = opentelemetry_value_to_json(value);
            serde_json::json!({
            "key": key.as_str(),
            "type": tpe,
            "value": value,
            })
        })
        .collect::<Vec<_>>();
    let references = if span.links.is_empty() {
        None
    } else {
        Some(
            span.links
                .iter()
                .map(|link| {
                    let span_context = &link.span_context;
                    serde_json::json!({
                        "refType": "FOLLOWS_FROM",
                        "traceID": span_context.trace_id().to_string(),
                        "spanID": span_context.span_id().to_string(),
                    })
                })
                .collect::<Vec<_>>(),
        )
    };
    serde_json::json!({
        "traceID": span.span_context.trace_id().to_string(),
        "spanID": span.span_context.span_id().to_string(),
        "startTime": span.start_time.duration_since(SystemTime::UNIX_EPOCH).expect("This does not fail").as_micros() as i64,
        "duration": span.end_time.duration_since(span.start_time).expect("This does not fail").as_micros() as i64,
        "operationName": span.name,
        "tags": tags,
        "logs": events,
        "flags": span.span_context.trace_flags().to_u8(),
        "processID": "p1",
        "warnings": None::<String>,
        "references": references,
    })
}

fn opentelemetry_value_to_json(value: &opentelemetry::Value) -> (&str, serde_json::Value) {
    match value {
        opentelemetry::Value::Bool(b) => ("bool", serde_json::json!(b)),
        opentelemetry::Value::I64(i) => ("int64", serde_json::json!(i)),
        opentelemetry::Value::F64(f) => ("float64", serde_json::json!(f)),
        opentelemetry::Value::String(s) => ("string", serde_json::json!(s.as_str())),
        v @ opentelemetry::Value::Array(_) => ("string", serde_json::json!(v.to_string())),
    }
}

/// Jaeger Json Runtime is an extension to [`TraceRuntime`].
///
/// [`TraceRuntime`]: opentelemetry::sdk::trace::TraceRuntime
#[async_trait]
pub trait JaegerJsonRuntime: TraceRuntime + std::fmt::Debug {
    /// Create a new directory if the given path does not exist yet
    async fn create_dir(&self, path: &Path) -> ExportResult;
    /// Write the provided content to a new file at the given path
    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult;
}

#[cfg(feature = "rt-tokio")]
#[async_trait]
impl JaegerJsonRuntime for opentelemetry::runtime::Tokio {
    async fn create_dir(&self, path: &Path) -> ExportResult {
        if tokio::fs::metadata(path).await.is_err() {
            tokio::fs::create_dir_all(path)
                .await
                .map_err(|e| TraceError::Other(Box::new(e)))?
        }

        Ok(())
    }

    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult {
        use tokio::io::AsyncWriteExt;

        let mut file = tokio::fs::File::create(path)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.write_all(content)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.sync_data()
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;

        Ok(())
    }
}

#[cfg(feature = "rt-tokio-current-thread")]
#[async_trait]
impl JaegerJsonRuntime for opentelemetry::runtime::TokioCurrentThread {
    async fn create_dir(&self, path: &Path) -> ExportResult {
        if tokio::fs::metadata(path).await.is_err() {
            tokio::fs::create_dir_all(path)
                .await
                .map_err(|e| TraceError::Other(Box::new(e)))?
        }

        Ok(())
    }

    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult {
        use tokio::io::AsyncWriteExt;

        let mut file = tokio::fs::File::create(path)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.write_all(content)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.sync_data()
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;

        Ok(())
    }
}

#[cfg(feature = "rt-async-std")]
#[async_trait]
impl JaegerJsonRuntime for opentelemetry::runtime::AsyncStd {
    async fn create_dir(&self, path: &Path) -> ExportResult {
        if async_std::fs::metadata(path).await.is_err() {
            async_std::fs::create_dir_all(path)
                .await
                .map_err(|e| TraceError::Other(Box::new(e)))?;
        }
        Ok(())
    }

    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult {
        use async_std::io::WriteExt;

        let mut file = async_std::fs::File::create(path)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.write_all(content)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.sync_data()
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;

        Ok(())
    }
}
<commit_msg>handle parent_span_id in jaeger_json (#907)<commit_after>//! # Jaeger JSON file Exporter
//!

use async_trait::async_trait;
use futures::{future::BoxFuture, FutureExt};
use opentelemetry::sdk::export::trace::{ExportResult, SpanData, SpanExporter};
use opentelemetry::sdk::trace::{TraceRuntime, Tracer};
use opentelemetry::trace::{SpanId, TraceError};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::SystemTime;

/// An exporter for jaeger comptible json files containing trace data
#[derive(Debug)]
pub struct JaegerJsonExporter<R> {
    out_path: PathBuf,
    file_prefix: String,
    service_name: String,
    runtime: R,
}

impl<R: JaegerJsonRuntime> JaegerJsonExporter<R> {
    /// Configure a new jaeger-json exporter
    ///
    /// * `out_path` refers to an directory where span data are written. If it does not exist, it is created by the exporter
    /// * `file_prefix` refers to a prefix prependend to each span file
    /// * `service_name` is used to identify the corresponding service in jaeger
    /// * `runtime` specifies the used async runtime to write the trace data
    pub fn new(out_path: PathBuf, file_prefix: String, service_name: String, runtime: R) -> Self {
        Self {
            out_path,
            file_prefix,
            service_name,
            runtime,
        }
    }

    /// Install the exporter using the internal provided runtime
    pub fn install_batch(self) -> Tracer {
        use opentelemetry::trace::TracerProvider;

        let runtime = self.runtime.clone();
        let provider_builder =
            opentelemetry::sdk::trace::TracerProvider::builder().with_batch_exporter(self, runtime);

        let provider = provider_builder.build();

        let tracer =
            provider.versioned_tracer("opentelemetry", Some(env!("CARGO_PKG_VERSION")), None);
        let _ = opentelemetry::global::set_tracer_provider(provider);

        tracer
    }
}

impl<R: JaegerJsonRuntime> SpanExporter for JaegerJsonExporter<R> {
    fn export(&mut self, batch: Vec<SpanData>) -> BoxFuture<'static, ExportResult> {
        let mut trace_map = HashMap::new();

        for span in batch {
            let ctx = &span.span_context;
            trace_map
                .entry(ctx.trace_id())
                .or_insert_with(Vec::new)
                .push(span_data_to_jaeger_json(span));
        }

        let data = trace_map
            .into_iter()
            .map(|(trace_id, spans)| {
                serde_json::json!({
                    "traceID": trace_id.to_string(),
                    "spans": spans,
                    "processes": {
                        "p1": {
                            "serviceName": self.service_name,
                            "tags": []
                        }
                    }
                })
            })
            .collect::<Vec<_>>();

        let json = serde_json::json!({
            "data": data,
        });

        let runtime = self.runtime.clone();
        let out_path = self.out_path.clone();
        let file_prefix = self.file_prefix.clone();

        async move {
            runtime.create_dir(&out_path).await?;

            let file_name = out_path.join(format!(
                "{}-{}.json",
                file_prefix,
                SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .expect("This does not fail")
                    .as_secs()
            ));
            runtime
                .write_to_file(
                    &file_name,
                    &serde_json::to_vec(&json).expect("This is a valid json value"),
                )
                .await?;

            Ok(())
        }
        .boxed()
    }
}

fn span_data_to_jaeger_json(
    span: opentelemetry::sdk::export::trace::SpanData,
) -> serde_json::Value {
    let events = span
        .events
        .iter()
        .map(|e| {
            let mut fields = e
                .attributes
                .iter()
                .map(|a| {
                    let (tpe, value) = opentelemetry_value_to_json(&a.value);
                    serde_json::json!({
                        "key": a.key.as_str(),
                        "type": tpe,
                        "value": value,
                    })
                })
                .collect::<Vec<_>>();
            fields.push(serde_json::json!({
                "key": "event",
                "type": "string",
                "value": e.name,
            }));

            serde_json::json!({
                "timestamp": e.timestamp.duration_since(SystemTime::UNIX_EPOCH).expect("This does not fail").as_micros() as i64,
                "fields": fields,
            })
        })
        .collect::<Vec<_>>();
    let tags = span
        .attributes
        .iter()
        .map(|(key, value)| {
            let (tpe, value) = opentelemetry_value_to_json(value);
            serde_json::json!({
            "key": key.as_str(),
            "type": tpe,
            "value": value,
            })
        })
        .collect::<Vec<_>>();
    let mut references = if span.links.is_empty() {
        None
    } else {
        Some(
            span.links
                .iter()
                .map(|link| {
                    let span_context = &link.span_context;
                    serde_json::json!({
                        "refType": "FOLLOWS_FROM",
                        "traceID": span_context.trace_id().to_string(),
                        "spanID": span_context.span_id().to_string(),
                    })
                })
                .collect::<Vec<_>>(),
        )
    };
    if span.parent_span_id != SpanId::INVALID {
        let val = serde_json::json!({
            "refType": "CHILD_OF",
            "traceID": span.span_context.trace_id().to_string(),
            "spanID": span.parent_span_id.to_string(),
        });
        references.get_or_insert_with(Vec::new).push(val);
    }
    serde_json::json!({
        "traceID": span.span_context.trace_id().to_string(),
        "spanID": span.span_context.span_id().to_string(),
        "startTime": span.start_time.duration_since(SystemTime::UNIX_EPOCH).expect("This does not fail").as_micros() as i64,
        "duration": span.end_time.duration_since(span.start_time).expect("This does not fail").as_micros() as i64,
        "operationName": span.name,
        "tags": tags,
        "logs": events,
        "flags": span.span_context.trace_flags().to_u8(),
        "processID": "p1",
        "warnings": None::<String>,
        "references": references,
    })
}

fn opentelemetry_value_to_json(value: &opentelemetry::Value) -> (&str, serde_json::Value) {
    match value {
        opentelemetry::Value::Bool(b) => ("bool", serde_json::json!(b)),
        opentelemetry::Value::I64(i) => ("int64", serde_json::json!(i)),
        opentelemetry::Value::F64(f) => ("float64", serde_json::json!(f)),
        opentelemetry::Value::String(s) => ("string", serde_json::json!(s.as_str())),
        v @ opentelemetry::Value::Array(_) => ("string", serde_json::json!(v.to_string())),
    }
}

/// Jaeger Json Runtime is an extension to [`TraceRuntime`].
///
/// [`TraceRuntime`]: opentelemetry::sdk::trace::TraceRuntime
#[async_trait]
pub trait JaegerJsonRuntime: TraceRuntime + std::fmt::Debug {
    /// Create a new directory if the given path does not exist yet
    async fn create_dir(&self, path: &Path) -> ExportResult;
    /// Write the provided content to a new file at the given path
    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult;
}

#[cfg(feature = "rt-tokio")]
#[async_trait]
impl JaegerJsonRuntime for opentelemetry::runtime::Tokio {
    async fn create_dir(&self, path: &Path) -> ExportResult {
        if tokio::fs::metadata(path).await.is_err() {
            tokio::fs::create_dir_all(path)
                .await
                .map_err(|e| TraceError::Other(Box::new(e)))?
        }

        Ok(())
    }

    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult {
        use tokio::io::AsyncWriteExt;

        let mut file = tokio::fs::File::create(path)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.write_all(content)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.sync_data()
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;

        Ok(())
    }
}

#[cfg(feature = "rt-tokio-current-thread")]
#[async_trait]
impl JaegerJsonRuntime for opentelemetry::runtime::TokioCurrentThread {
    async fn create_dir(&self, path: &Path) -> ExportResult {
        if tokio::fs::metadata(path).await.is_err() {
            tokio::fs::create_dir_all(path)
                .await
                .map_err(|e| TraceError::Other(Box::new(e)))?
        }

        Ok(())
    }

    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult {
        use tokio::io::AsyncWriteExt;

        let mut file = tokio::fs::File::create(path)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.write_all(content)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.sync_data()
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;

        Ok(())
    }
}

#[cfg(feature = "rt-async-std")]
#[async_trait]
impl JaegerJsonRuntime for opentelemetry::runtime::AsyncStd {
    async fn create_dir(&self, path: &Path) -> ExportResult {
        if async_std::fs::metadata(path).await.is_err() {
            async_std::fs::create_dir_all(path)
                .await
                .map_err(|e| TraceError::Other(Box::new(e)))?;
        }
        Ok(())
    }

    async fn write_to_file(&self, path: &Path, content: &[u8]) -> ExportResult {
        use async_std::io::WriteExt;

        let mut file = async_std::fs::File::create(path)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.write_all(content)
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;
        file.sync_data()
            .await
            .map_err(|e| TraceError::Other(Box::new(e)))?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::io::IoError;

pub type ProtobufResult<T> = Result<T, ProtobufError>;

#[deriving(Show,Eq,PartialEq)]
pub enum ProtobufError {
    IoError(IoError),
    WireError(String),
}
<commit_msg>Implement Error for ProtobufError<commit_after>use std::io::IoError;
use std::error::Error;

pub type ProtobufResult<T> = Result<T, ProtobufError>;

#[deriving(Show,Eq,PartialEq)]
pub enum ProtobufError {
    IoError(IoError),
    WireError(String),
}

impl Error for ProtobufError {
    fn description(&self) -> &str {
        match self {
            // not sure that cause should be included in message
            &ProtobufError::IoError(ref e) => e.description(),
            &ProtobufError::WireError(ref e) => e.as_slice(),
        }
    }

    fn cause(&self) -> Option<&Error> {
        match self {
            &ProtobufError::IoError(ref e) => Some(e as &Error),
            &ProtobufError::WireError(..) => None,
        }
    }
}
<|endoftext|>


<commit_before>use ra_parser::{TokenSource, Token};
use ra_syntax::{classify_literal, SmolStr, SyntaxKind, SyntaxKind::*, T};
use std::cell::{RefCell, Cell};
use tt::buffer::{TokenBuffer, Cursor};

#[derive(Debug, Clone, Eq, PartialEq)]
struct TtToken {
    pub kind: SyntaxKind,
    pub is_joint_to_next: bool,
    pub text: SmolStr,
}

pub(crate) struct SubtreeTokenSource<'a> {
    cached_cursor: Cell<Cursor<'a>>,
    cached: RefCell<Vec<Option<TtToken>>>,
    curr: (Token, usize),
}

impl<'a> SubtreeTokenSource<'a> {
    // Helper function used in test
    #[allow(unused)]
    pub fn text(&self) -> SmolStr {
        match self.get(self.curr.1) {
            Some(tt) => tt.text,
            _ => SmolStr::new(""),
        }
    }
}

impl<'a> SubtreeTokenSource<'a> {
    pub fn new(buffer: &'a TokenBuffer) -> SubtreeTokenSource<'a> {
        let cursor = buffer.begin();

        let mut res = SubtreeTokenSource {
            curr: (Token { kind: EOF, is_jointed_to_next: false }, 0),
            cached_cursor: Cell::new(cursor),
            cached: RefCell::new(Vec::with_capacity(10)),
        };
        res.curr = (res.mk_token(0), 0);
        res
    }

    fn mk_token(&self, pos: usize) -> Token {
        match self.get(pos) {
            Some(tt) => Token { kind: tt.kind, is_jointed_to_next: tt.is_joint_to_next },
            None => Token { kind: EOF, is_jointed_to_next: false },
        }
    }

    fn get(&self, pos: usize) -> Option<TtToken> {
        let mut cached = self.cached.borrow_mut();
        if pos < cached.len() {
            return cached[pos].clone();
        }

        while pos >= cached.len() {
            let cursor = self.cached_cursor.get();
            if cursor.eof() {
                cached.push(None);
                continue;
            }

            match cursor.token_tree() {
                Some(tt::TokenTree::Leaf(leaf)) => {
                    cached.push(Some(convert_leaf(&leaf)));
                    self.cached_cursor.set(cursor.bump());
                }
                Some(tt::TokenTree::Subtree(subtree)) => {
                    self.cached_cursor.set(cursor.subtree().unwrap());
                    cached.push(Some(convert_delim(subtree.delimiter, false)));
                }
                None => {
                    if let Some(subtree) = cursor.end() {
                        cached.push(Some(convert_delim(subtree.delimiter, true)));
                        self.cached_cursor.set(cursor.bump());
                    }
                }
            }
        }

        return cached[pos].clone();
    }
}

impl<'a> TokenSource for SubtreeTokenSource<'a> {
    fn current(&self) -> Token {
        self.curr.0
    }

    /// Lookahead n token
    fn lookahead_nth(&self, n: usize) -> Token {
        self.mk_token(self.curr.1 + n)
    }

    /// bump cursor to next token
    fn bump(&mut self) {
        if self.current().kind == EOF {
            return;
        }

        self.curr = (self.mk_token(self.curr.1 + 1), self.curr.1 + 1);
    }

    /// Is the current token a specified keyword?
    fn is_keyword(&self, kw: &str) -> bool {
        match self.get(self.curr.1) {
            Some(t) => t.text == *kw,
            _ => false,
        }
    }
}

fn convert_delim(d: tt::Delimiter, closing: bool) -> TtToken {
    let (kinds, texts) = match d {
        tt::Delimiter::Parenthesis => ([T!['('], T![')']], "()"),
        tt::Delimiter::Brace => ([T!['{'], T!['}']], "{}"),
        tt::Delimiter::Bracket => ([T!['['], T![']']], "[]"),
        tt::Delimiter::None => ([L_DOLLAR, R_DOLLAR], ""),
    };

    let idx = closing as usize;
    let kind = kinds[idx];
    let text = if texts.len() > 0 { &texts[idx..texts.len() - (1 - idx)] } else { "" };
    TtToken { kind, is_joint_to_next: false, text: SmolStr::new(text) }
}

fn convert_literal(l: &tt::Literal) -> TtToken {
    let kind =
        classify_literal(&l.text).map(|tkn| tkn.kind).unwrap_or_else(|| match l.text.as_ref() {
            "true" => T![true],
            "false" => T![false],
            _ => panic!("Fail to convert given literal {:#?}", &l),
        });

    TtToken { kind, is_joint_to_next: false, text: l.text.clone() }
}

fn convert_ident(ident: &tt::Ident) -> TtToken {
    let kind = if let Some('\'') = ident.text.chars().next() {
        LIFETIME
    } else {
        SyntaxKind::from_keyword(ident.text.as_str()).unwrap_or(IDENT)
    };

    TtToken { kind, is_joint_to_next: false, text: ident.text.clone() }
}

fn convert_punct(p: &tt::Punct) -> TtToken {
    let kind = match p.char {
        // lexer may produce compound tokens for these ones
        '.' => T![.],
        ':' => T![:],
        '=' => T![=],
        '!' => T![!],
        '-' => T![-],
        c => SyntaxKind::from_char(c).unwrap(),
    };
    let text = {
        let mut buf = [0u8; 4];
        let s: &str = p.char.encode_utf8(&mut buf);
        SmolStr::new(s)
    };
    TtToken { kind, is_joint_to_next: p.spacing == tt::Spacing::Joint, text }
}

fn convert_leaf(leaf: &tt::Leaf) -> TtToken {
    match leaf {
        tt::Leaf::Literal(l) => convert_literal(l),
        tt::Leaf::Ident(ident) => convert_ident(ident),
        tt::Leaf::Punct(punct) => convert_punct(punct),
    }
}
<commit_msg>Use cfg(test) instead of allow(unused)<commit_after>use ra_parser::{TokenSource, Token};
use ra_syntax::{classify_literal, SmolStr, SyntaxKind, SyntaxKind::*, T};
use std::cell::{RefCell, Cell};
use tt::buffer::{TokenBuffer, Cursor};

#[derive(Debug, Clone, Eq, PartialEq)]
struct TtToken {
    pub kind: SyntaxKind,
    pub is_joint_to_next: bool,
    pub text: SmolStr,
}

pub(crate) struct SubtreeTokenSource<'a> {
    cached_cursor: Cell<Cursor<'a>>,
    cached: RefCell<Vec<Option<TtToken>>>,
    curr: (Token, usize),
}

impl<'a> SubtreeTokenSource<'a> {
    // Helper function used in test
    #[cfg(test)]
    pub fn text(&self) -> SmolStr {
        match self.get(self.curr.1) {
            Some(tt) => tt.text,
            _ => SmolStr::new(""),
        }
    }
}

impl<'a> SubtreeTokenSource<'a> {
    pub fn new(buffer: &'a TokenBuffer) -> SubtreeTokenSource<'a> {
        let cursor = buffer.begin();

        let mut res = SubtreeTokenSource {
            curr: (Token { kind: EOF, is_jointed_to_next: false }, 0),
            cached_cursor: Cell::new(cursor),
            cached: RefCell::new(Vec::with_capacity(10)),
        };
        res.curr = (res.mk_token(0), 0);
        res
    }

    fn mk_token(&self, pos: usize) -> Token {
        match self.get(pos) {
            Some(tt) => Token { kind: tt.kind, is_jointed_to_next: tt.is_joint_to_next },
            None => Token { kind: EOF, is_jointed_to_next: false },
        }
    }

    fn get(&self, pos: usize) -> Option<TtToken> {
        let mut cached = self.cached.borrow_mut();
        if pos < cached.len() {
            return cached[pos].clone();
        }

        while pos >= cached.len() {
            let cursor = self.cached_cursor.get();
            if cursor.eof() {
                cached.push(None);
                continue;
            }

            match cursor.token_tree() {
                Some(tt::TokenTree::Leaf(leaf)) => {
                    cached.push(Some(convert_leaf(&leaf)));
                    self.cached_cursor.set(cursor.bump());
                }
                Some(tt::TokenTree::Subtree(subtree)) => {
                    self.cached_cursor.set(cursor.subtree().unwrap());
                    cached.push(Some(convert_delim(subtree.delimiter, false)));
                }
                None => {
                    if let Some(subtree) = cursor.end() {
                        cached.push(Some(convert_delim(subtree.delimiter, true)));
                        self.cached_cursor.set(cursor.bump());
                    }
                }
            }
        }

        return cached[pos].clone();
    }
}

impl<'a> TokenSource for SubtreeTokenSource<'a> {
    fn current(&self) -> Token {
        self.curr.0
    }

    /// Lookahead n token
    fn lookahead_nth(&self, n: usize) -> Token {
        self.mk_token(self.curr.1 + n)
    }

    /// bump cursor to next token
    fn bump(&mut self) {
        if self.current().kind == EOF {
            return;
        }

        self.curr = (self.mk_token(self.curr.1 + 1), self.curr.1 + 1);
    }

    /// Is the current token a specified keyword?
    fn is_keyword(&self, kw: &str) -> bool {
        match self.get(self.curr.1) {
            Some(t) => t.text == *kw,
            _ => false,
        }
    }
}

fn convert_delim(d: tt::Delimiter, closing: bool) -> TtToken {
    let (kinds, texts) = match d {
        tt::Delimiter::Parenthesis => ([T!['('], T![')']], "()"),
        tt::Delimiter::Brace => ([T!['{'], T!['}']], "{}"),
        tt::Delimiter::Bracket => ([T!['['], T![']']], "[]"),
        tt::Delimiter::None => ([L_DOLLAR, R_DOLLAR], ""),
    };

    let idx = closing as usize;
    let kind = kinds[idx];
    let text = if texts.len() > 0 { &texts[idx..texts.len() - (1 - idx)] } else { "" };
    TtToken { kind, is_joint_to_next: false, text: SmolStr::new(text) }
}

fn convert_literal(l: &tt::Literal) -> TtToken {
    let kind =
        classify_literal(&l.text).map(|tkn| tkn.kind).unwrap_or_else(|| match l.text.as_ref() {
            "true" => T![true],
            "false" => T![false],
            _ => panic!("Fail to convert given literal {:#?}", &l),
        });

    TtToken { kind, is_joint_to_next: false, text: l.text.clone() }
}

fn convert_ident(ident: &tt::Ident) -> TtToken {
    let kind = if let Some('\'') = ident.text.chars().next() {
        LIFETIME
    } else {
        SyntaxKind::from_keyword(ident.text.as_str()).unwrap_or(IDENT)
    };

    TtToken { kind, is_joint_to_next: false, text: ident.text.clone() }
}

fn convert_punct(p: &tt::Punct) -> TtToken {
    let kind = match p.char {
        // lexer may produce compound tokens for these ones
        '.' => T![.],
        ':' => T![:],
        '=' => T![=],
        '!' => T![!],
        '-' => T![-],
        c => SyntaxKind::from_char(c).unwrap(),
    };
    let text = {
        let mut buf = [0u8; 4];
        let s: &str = p.char.encode_utf8(&mut buf);
        SmolStr::new(s)
    };
    TtToken { kind, is_joint_to_next: p.spacing == tt::Spacing::Joint, text }
}

fn convert_leaf(leaf: &tt::Leaf) -> TtToken {
    match leaf {
        tt::Leaf::Literal(l) => convert_literal(l),
        tt::Leaf::Ident(ident) => convert_ident(ident),
        tt::Leaf::Punct(punct) => convert_punct(punct),
    }
}
<|endoftext|>


<commit_before>#![feature(stdsimd)]
#![feature(stdsimd_internal)]

extern crate std_detect;

use std_detect::detect::{features, Feature};

#[test]
fn features_roundtrip() {
    for (f, _) in features() {
        let _ = Feature::from_str(f).unwrap();
    }
}
<commit_msg>Drop the features test for now<commit_after><|endoftext|>


<commit_before>/// The waker registration is derived from smoltcp waker registration
use core::task::Waker;

/// Utility struct to register and wake a waker.
#[derive(Debug)]
pub(crate) struct WakerRegistration {
	waker: Option<Waker>,
}

impl WakerRegistration {
	pub(crate) const fn new() -> Self {
		Self { waker: None }
	}

	/// Register a waker. Overwrites the previous waker, if any.
	pub(crate) fn register(&mut self, w: &Waker) {
		match self.waker {
			// Optimization: If both the old and new Wakers wake the same task, we can simply
			// keep the old waker, skipping the clone. (In most executor implementations,
			// cloning a waker is somewhat expensive, comparable to cloning an Arc).
			Some(ref w2) if (w2.will_wake(w)) => {}
			_ => {
				// clone the new waker and store it
				if let Some(old_waker) = std::mem::replace(&mut self.waker, Some(w.clone())) {
					// We had a waker registered for another task. Wake it, so the other task can
					// reregister itself if it's still interested.
					//
					// If two tasks are waiting on the same thing concurrently, this will cause them
					// to wake each other in a loop fighting over this WakerRegistration. This wastes
					// CPU but things will still work.
					//
					// If the user wants to have two tasks waiting on the same thing they should use
					// a more appropriate primitive that can store multiple wakers.
					old_waker.wake()
				}
			}
		}
	}

	/// Wake the registered waker, if any.
	pub(crate) fn wake(&mut self) {
		self.waker.take().map(|w| w.wake());
	}
}
<commit_msg>remove clippy warnings<commit_after>/// The waker registration is derived from smoltcp waker registration
use core::task::Waker;

/// Utility struct to register and wake a waker.
#[derive(Debug)]
pub(crate) struct WakerRegistration {
	waker: Option<Waker>,
}

impl WakerRegistration {
	pub(crate) const fn new() -> Self {
		Self { waker: None }
	}

	/// Register a waker. Overwrites the previous waker, if any.
	pub(crate) fn register(&mut self, w: &Waker) {
		match self.waker {
			// Optimization: If both the old and new Wakers wake the same task, we can simply
			// keep the old waker, skipping the clone. (In most executor implementations,
			// cloning a waker is somewhat expensive, comparable to cloning an Arc).
			Some(ref w2) if (w2.will_wake(w)) => {}
			_ => {
				// clone the new waker and store it
				if let Some(old_waker) = std::mem::replace(&mut self.waker, Some(w.clone())) {
					// We had a waker registered for another task. Wake it, so the other task can
					// reregister itself if it's still interested.
					//
					// If two tasks are waiting on the same thing concurrently, this will cause them
					// to wake each other in a loop fighting over this WakerRegistration. This wastes
					// CPU but things will still work.
					//
					// If the user wants to have two tasks waiting on the same thing they should use
					// a more appropriate primitive that can store multiple wakers.
					old_waker.wake()
				}
			}
		}
	}

	/// Wake the registered waker, if any.
	pub(crate) fn wake(&mut self) {
		if let Some(w) = self.waker.take() {
			w.wake()
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use super::xhci_abi::{
    AddressedTrb, Error as TrbError, LinkTrb, TransferDescriptor, Trb, TrbCast, TrbType,
};
use std::fmt::{self, Display};
use std::mem::size_of;
use sys_util::{GuestAddress, GuestMemory, GuestMemoryError};

#[derive(Debug)]
pub enum Error {
    ReadGuestMemory(GuestMemoryError),
    BadDequeuePointer(GuestAddress),
    CastTrb(TrbError),
    TrbChain(TrbError),
}

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

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

        match self {
            ReadGuestMemory(e) => write!(f, "cannot read guest memory: {}", e),
            BadDequeuePointer(addr) => write!(f, "bad dequeue pointer: {}", addr),
            CastTrb(e) => write!(f, "cannot cast trb: {}", e),
            TrbChain(e) => write!(f, "cannot get trb chain bit: {}", e),
        }
    }
}

/// Ring Buffer is segmented circular buffer in guest memory containing work items
/// called transfer descriptors, each of which consists of one or more TRBs.
/// Ring buffer logic is shared between transfer ring and command ring.
/// Transfer Ring management is defined in xHCI spec 4.9.2.
pub struct RingBuffer {
    name: String,
    mem: GuestMemory,
    dequeue_pointer: GuestAddress,
    // Used to check if the ring is empty. Toggled when looping back to the begining
    // of the buffer.
    consumer_cycle_state: bool,
}

impl Display for RingBuffer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "RingBuffer `{}`", self.name)
    }
}

// Public interfaces for Ring buffer.
impl RingBuffer {
    /// Create a new RingBuffer.
    pub fn new(name: String, mem: GuestMemory) -> Self {
        RingBuffer {
            name,
            mem,
            dequeue_pointer: GuestAddress(0),
            consumer_cycle_state: false,
        }
    }

    /// Dequeue next transfer descriptor from the transfer ring.
    pub fn dequeue_transfer_descriptor(&mut self) -> Result<Option<TransferDescriptor>> {
        let mut td: TransferDescriptor = TransferDescriptor::new();
        while let Some(addressed_trb) = self.get_current_trb()? {
            if let Ok(TrbType::Link) = addressed_trb.trb.get_trb_type() {
                let link_trb = addressed_trb
                    .trb
                    .cast::<LinkTrb>()
                    .map_err(Error::CastTrb)?;
                self.dequeue_pointer = GuestAddress(link_trb.get_ring_segment_pointer());
                self.consumer_cycle_state =
                    self.consumer_cycle_state != link_trb.get_toggle_cycle();
                continue;
            }

            self.dequeue_pointer = match self.dequeue_pointer.checked_add(size_of::<Trb>() as u64) {
                Some(addr) => addr,
                None => {
                    return Err(Error::BadDequeuePointer(self.dequeue_pointer));
                }
            };

            usb_debug!(
                "{}: adding trb to td {}",
                self.name.as_str(),
                addressed_trb.trb
            );
            td.push(addressed_trb);
            if !addressed_trb.trb.get_chain_bit().map_err(Error::TrbChain)? {
                usb_debug!("trb chain is false returning");
                break;
            }
        }
        // A valid transfer descriptor contains at least one addressed trb and the last trb has
        // chain bit != 0.
        match td.last() {
            Some(t) => {
                if t.trb.get_chain_bit().map_err(Error::TrbChain)? {
                    return Ok(None);
                }
            }
            None => return Ok(None),
        }
        Ok(Some(td))
    }

    /// Set dequeue pointer of the ring buffer.
    pub fn set_dequeue_pointer(&mut self, addr: GuestAddress) {
        usb_debug!("{}: set dequeue pointer {:x}", self.name.as_str(), addr.0);

        self.dequeue_pointer = addr;
    }

    /// Set consumer cycle state of the ring buffer.
    pub fn set_consumer_cycle_state(&mut self, state: bool) {
        usb_debug!("{}: set consumer cycle state {}", self.name.as_str(), state);
        self.consumer_cycle_state = state;
    }

    // Read trb pointed by dequeue pointer. Does not proceed dequeue pointer.
    fn get_current_trb(&self) -> Result<Option<AddressedTrb>> {
        let trb: Trb = self
            .mem
            .read_obj_from_addr(self.dequeue_pointer)
            .map_err(Error::ReadGuestMemory)?;
        usb_debug!("{}: trb read from memory {:?}", self.name.as_str(), trb);
        // If cycle bit of trb does not equal consumer cycle state, the ring is empty.
        // This trb is invalid.
        if trb.get_cycle() != self.consumer_cycle_state {
            usb_debug!(
                "cycle bit does not match, self cycle {}",
                self.consumer_cycle_state
            );
            Ok(None)
        } else {
            Ok(Some(AddressedTrb {
                trb,
                gpa: self.dequeue_pointer.0,
            }))
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::usb::xhci::xhci_abi::*;

    #[test]
    fn ring_test_dequeue() {
        let trb_size = size_of::<Trb>() as u64;
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
        let mut transfer_ring = RingBuffer::new(String::new(), gm.clone());

        // Structure of ring buffer:
        //  0x100  --> 0x200  --> 0x300
        //  trb 1  |   trb 3  |   trb 5
        //  trb 2  |   trb 4  |   trb 6
        //  l trb  -   l trb  -   l trb to 0x100
        let mut trb = NormalTrb::new();
        trb.set_trb_type(TrbType::Normal);
        trb.set_data_buffer(1);
        trb.set_chain(true);
        gm.write_obj_at_addr(trb.clone(), GuestAddress(0x100))
            .unwrap();

        trb.set_data_buffer(2);
        gm.write_obj_at_addr(trb, GuestAddress(0x100 + trb_size))
            .unwrap();

        let mut ltrb = LinkTrb::new();
        ltrb.set_trb_type(TrbType::Link);
        ltrb.set_ring_segment_pointer(0x200);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x100 + 2 * trb_size))
            .unwrap();

        trb.set_data_buffer(3);
        gm.write_obj_at_addr(trb, GuestAddress(0x200)).unwrap();

        // Chain bit is false.
        trb.set_data_buffer(4);
        trb.set_chain(false);
        gm.write_obj_at_addr(trb, GuestAddress(0x200 + 1 * trb_size))
            .unwrap();

        ltrb.set_ring_segment_pointer(0x300);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x200 + 2 * trb_size))
            .unwrap();

        trb.set_data_buffer(5);
        trb.set_chain(true);
        gm.write_obj_at_addr(trb, GuestAddress(0x300)).unwrap();

        // Chain bit is false.
        trb.set_data_buffer(6);
        trb.set_chain(false);
        gm.write_obj_at_addr(trb, GuestAddress(0x300 + 1 * trb_size))
            .unwrap();

        ltrb.set_ring_segment_pointer(0x100);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x300 + 2 * trb_size))
            .unwrap();

        transfer_ring.set_dequeue_pointer(GuestAddress(0x100));
        transfer_ring.set_consumer_cycle_state(false);

        // Read first transfer descriptor.
        let descriptor = transfer_ring
            .dequeue_transfer_descriptor()
            .unwrap()
            .unwrap();
        assert_eq!(descriptor.len(), 4);
        assert_eq!(descriptor[0].trb.get_parameter(), 1);
        assert_eq!(descriptor[1].trb.get_parameter(), 2);
        assert_eq!(descriptor[2].trb.get_parameter(), 3);
        assert_eq!(descriptor[3].trb.get_parameter(), 4);

        // Read second transfer descriptor.
        let descriptor = transfer_ring
            .dequeue_transfer_descriptor()
            .unwrap()
            .unwrap();
        assert_eq!(descriptor.len(), 2);
        assert_eq!(descriptor[0].trb.get_parameter(), 5);
        assert_eq!(descriptor[1].trb.get_parameter(), 6);
    }

    #[test]
    fn transfer_ring_test_dequeue_failure() {
        let trb_size = size_of::<Trb>() as u64;
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
        let mut transfer_ring = RingBuffer::new(String::new(), gm.clone());

        let mut trb = NormalTrb::new();
        trb.set_trb_type(TrbType::Normal);
        trb.set_data_buffer(1);
        trb.set_chain(true);
        gm.write_obj_at_addr(trb.clone(), GuestAddress(0x100))
            .unwrap();

        trb.set_data_buffer(2);
        gm.write_obj_at_addr(trb, GuestAddress(0x100 + trb_size))
            .unwrap();

        let mut ltrb = LinkTrb::new();
        ltrb.set_trb_type(TrbType::Link);
        ltrb.set_ring_segment_pointer(0x200);
        ltrb.set_toggle_cycle(true);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x100 + 2 * trb_size))
            .unwrap();

        trb.set_data_buffer(3);
        gm.write_obj_at_addr(trb, GuestAddress(0x200)).unwrap();

        transfer_ring.set_dequeue_pointer(GuestAddress(0x100));
        transfer_ring.set_consumer_cycle_state(false);

        // Read first transfer descriptor.
        let descriptor = transfer_ring.dequeue_transfer_descriptor().unwrap();
        assert_eq!(descriptor.is_none(), true);
    }
}
<commit_msg>devices: usb: add unit test for ring buffer cycle<commit_after>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use super::xhci_abi::{
    AddressedTrb, Error as TrbError, LinkTrb, TransferDescriptor, Trb, TrbCast, TrbType,
};
use std::fmt::{self, Display};
use std::mem::size_of;
use sys_util::{GuestAddress, GuestMemory, GuestMemoryError};

#[derive(Debug)]
pub enum Error {
    ReadGuestMemory(GuestMemoryError),
    BadDequeuePointer(GuestAddress),
    CastTrb(TrbError),
    TrbChain(TrbError),
}

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

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

        match self {
            ReadGuestMemory(e) => write!(f, "cannot read guest memory: {}", e),
            BadDequeuePointer(addr) => write!(f, "bad dequeue pointer: {}", addr),
            CastTrb(e) => write!(f, "cannot cast trb: {}", e),
            TrbChain(e) => write!(f, "cannot get trb chain bit: {}", e),
        }
    }
}

/// Ring Buffer is segmented circular buffer in guest memory containing work items
/// called transfer descriptors, each of which consists of one or more TRBs.
/// Ring buffer logic is shared between transfer ring and command ring.
/// Transfer Ring management is defined in xHCI spec 4.9.2.
pub struct RingBuffer {
    name: String,
    mem: GuestMemory,
    dequeue_pointer: GuestAddress,
    // Used to check if the ring is empty. Toggled when looping back to the begining
    // of the buffer.
    consumer_cycle_state: bool,
}

impl Display for RingBuffer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "RingBuffer `{}`", self.name)
    }
}

// Public interfaces for Ring buffer.
impl RingBuffer {
    /// Create a new RingBuffer.
    pub fn new(name: String, mem: GuestMemory) -> Self {
        RingBuffer {
            name,
            mem,
            dequeue_pointer: GuestAddress(0),
            consumer_cycle_state: false,
        }
    }

    /// Dequeue next transfer descriptor from the transfer ring.
    pub fn dequeue_transfer_descriptor(&mut self) -> Result<Option<TransferDescriptor>> {
        let mut td: TransferDescriptor = TransferDescriptor::new();
        while let Some(addressed_trb) = self.get_current_trb()? {
            if let Ok(TrbType::Link) = addressed_trb.trb.get_trb_type() {
                let link_trb = addressed_trb
                    .trb
                    .cast::<LinkTrb>()
                    .map_err(Error::CastTrb)?;
                self.dequeue_pointer = GuestAddress(link_trb.get_ring_segment_pointer());
                self.consumer_cycle_state =
                    self.consumer_cycle_state != link_trb.get_toggle_cycle();
                continue;
            }

            self.dequeue_pointer = match self.dequeue_pointer.checked_add(size_of::<Trb>() as u64) {
                Some(addr) => addr,
                None => {
                    return Err(Error::BadDequeuePointer(self.dequeue_pointer));
                }
            };

            usb_debug!(
                "{}: adding trb to td {}",
                self.name.as_str(),
                addressed_trb.trb
            );
            td.push(addressed_trb);
            if !addressed_trb.trb.get_chain_bit().map_err(Error::TrbChain)? {
                usb_debug!("trb chain is false returning");
                break;
            }
        }
        // A valid transfer descriptor contains at least one addressed trb and the last trb has
        // chain bit != 0.
        match td.last() {
            Some(t) => {
                if t.trb.get_chain_bit().map_err(Error::TrbChain)? {
                    return Ok(None);
                }
            }
            None => return Ok(None),
        }
        Ok(Some(td))
    }

    /// Set dequeue pointer of the ring buffer.
    pub fn set_dequeue_pointer(&mut self, addr: GuestAddress) {
        usb_debug!("{}: set dequeue pointer {:x}", self.name.as_str(), addr.0);

        self.dequeue_pointer = addr;
    }

    /// Set consumer cycle state of the ring buffer.
    pub fn set_consumer_cycle_state(&mut self, state: bool) {
        usb_debug!("{}: set consumer cycle state {}", self.name.as_str(), state);
        self.consumer_cycle_state = state;
    }

    // Read trb pointed by dequeue pointer. Does not proceed dequeue pointer.
    fn get_current_trb(&self) -> Result<Option<AddressedTrb>> {
        let trb: Trb = self
            .mem
            .read_obj_from_addr(self.dequeue_pointer)
            .map_err(Error::ReadGuestMemory)?;
        usb_debug!("{}: trb read from memory {:?}", self.name.as_str(), trb);
        // If cycle bit of trb does not equal consumer cycle state, the ring is empty.
        // This trb is invalid.
        if trb.get_cycle() != self.consumer_cycle_state {
            usb_debug!(
                "cycle bit does not match, self cycle {}",
                self.consumer_cycle_state
            );
            Ok(None)
        } else {
            Ok(Some(AddressedTrb {
                trb,
                gpa: self.dequeue_pointer.0,
            }))
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::usb::xhci::xhci_abi::*;

    #[test]
    fn ring_test_dequeue() {
        let trb_size = size_of::<Trb>() as u64;
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
        let mut transfer_ring = RingBuffer::new(String::new(), gm.clone());

        // Structure of ring buffer:
        //  0x100  --> 0x200  --> 0x300
        //  trb 1  |   trb 3  |   trb 5
        //  trb 2  |   trb 4  |   trb 6
        //  l trb  -   l trb  -   l trb to 0x100
        let mut trb = NormalTrb::new();
        trb.set_trb_type(TrbType::Normal);
        trb.set_data_buffer(1);
        trb.set_chain(true);
        gm.write_obj_at_addr(trb.clone(), GuestAddress(0x100))
            .unwrap();

        trb.set_data_buffer(2);
        gm.write_obj_at_addr(trb, GuestAddress(0x100 + trb_size))
            .unwrap();

        let mut ltrb = LinkTrb::new();
        ltrb.set_trb_type(TrbType::Link);
        ltrb.set_ring_segment_pointer(0x200);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x100 + 2 * trb_size))
            .unwrap();

        trb.set_data_buffer(3);
        gm.write_obj_at_addr(trb, GuestAddress(0x200)).unwrap();

        // Chain bit is false.
        trb.set_data_buffer(4);
        trb.set_chain(false);
        gm.write_obj_at_addr(trb, GuestAddress(0x200 + 1 * trb_size))
            .unwrap();

        ltrb.set_ring_segment_pointer(0x300);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x200 + 2 * trb_size))
            .unwrap();

        trb.set_data_buffer(5);
        trb.set_chain(true);
        gm.write_obj_at_addr(trb, GuestAddress(0x300)).unwrap();

        // Chain bit is false.
        trb.set_data_buffer(6);
        trb.set_chain(false);
        gm.write_obj_at_addr(trb, GuestAddress(0x300 + 1 * trb_size))
            .unwrap();

        ltrb.set_ring_segment_pointer(0x100);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x300 + 2 * trb_size))
            .unwrap();

        transfer_ring.set_dequeue_pointer(GuestAddress(0x100));
        transfer_ring.set_consumer_cycle_state(false);

        // Read first transfer descriptor.
        let descriptor = transfer_ring
            .dequeue_transfer_descriptor()
            .unwrap()
            .unwrap();
        assert_eq!(descriptor.len(), 4);
        assert_eq!(descriptor[0].trb.get_parameter(), 1);
        assert_eq!(descriptor[1].trb.get_parameter(), 2);
        assert_eq!(descriptor[2].trb.get_parameter(), 3);
        assert_eq!(descriptor[3].trb.get_parameter(), 4);

        // Read second transfer descriptor.
        let descriptor = transfer_ring
            .dequeue_transfer_descriptor()
            .unwrap()
            .unwrap();
        assert_eq!(descriptor.len(), 2);
        assert_eq!(descriptor[0].trb.get_parameter(), 5);
        assert_eq!(descriptor[1].trb.get_parameter(), 6);
    }

    #[test]
    fn transfer_ring_test_dequeue_failure() {
        let trb_size = size_of::<Trb>() as u64;
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
        let mut transfer_ring = RingBuffer::new(String::new(), gm.clone());

        let mut trb = NormalTrb::new();
        trb.set_trb_type(TrbType::Normal);
        trb.set_data_buffer(1);
        trb.set_chain(true);
        gm.write_obj_at_addr(trb.clone(), GuestAddress(0x100))
            .unwrap();

        trb.set_data_buffer(2);
        gm.write_obj_at_addr(trb, GuestAddress(0x100 + trb_size))
            .unwrap();

        let mut ltrb = LinkTrb::new();
        ltrb.set_trb_type(TrbType::Link);
        ltrb.set_ring_segment_pointer(0x200);
        ltrb.set_toggle_cycle(true);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x100 + 2 * trb_size))
            .unwrap();

        trb.set_data_buffer(3);
        gm.write_obj_at_addr(trb, GuestAddress(0x200)).unwrap();

        transfer_ring.set_dequeue_pointer(GuestAddress(0x100));
        transfer_ring.set_consumer_cycle_state(false);

        // Read first transfer descriptor.
        let descriptor = transfer_ring.dequeue_transfer_descriptor().unwrap();
        assert_eq!(descriptor.is_none(), true);
    }

    #[test]
    fn ring_test_toggle_cycle() {
        let trb_size = size_of::<Trb>() as u64;
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
        let mut transfer_ring = RingBuffer::new(String::new(), gm.clone());

        let mut trb = NormalTrb::new();
        trb.set_trb_type(TrbType::Normal);
        trb.set_data_buffer(1);
        trb.set_chain(false);
        trb.set_cycle(false);
        gm.write_obj_at_addr(trb.clone(), GuestAddress(0x100))
            .unwrap();

        let mut ltrb = LinkTrb::new();
        ltrb.set_trb_type(TrbType::Link);
        ltrb.set_ring_segment_pointer(0x100);
        ltrb.set_toggle_cycle(true);
        ltrb.set_cycle(false);
        gm.write_obj_at_addr(ltrb, GuestAddress(0x100 + trb_size))
            .unwrap();

        // Initial state: consumer cycle = false
        transfer_ring.set_dequeue_pointer(GuestAddress(0x100));
        transfer_ring.set_consumer_cycle_state(false);

        // Read first transfer descriptor.
        let descriptor = transfer_ring
            .dequeue_transfer_descriptor()
            .unwrap()
            .unwrap();
        assert_eq!(descriptor.len(), 1);
        assert_eq!(descriptor[0].trb.get_parameter(), 1);

        // Cycle bit should be unchanged since we haven't advanced past the Link TRB yet.
        assert_eq!(transfer_ring.consumer_cycle_state, false);

        // Overwrite the first TRB with a new one (data = 2)
        // with the new producer cycle bit state (true).
        let mut trb = NormalTrb::new();
        trb.set_trb_type(TrbType::Normal);
        trb.set_data_buffer(2);
        trb.set_cycle(true); // Link TRB toggled the cycle.
        gm.write_obj_at_addr(trb.clone(), GuestAddress(0x100))
            .unwrap();

        // Read new transfer descriptor.
        let descriptor = transfer_ring
            .dequeue_transfer_descriptor()
            .unwrap()
            .unwrap();
        assert_eq!(descriptor.len(), 1);
        assert_eq!(descriptor[0].trb.get_parameter(), 2);

        assert_eq!(transfer_ring.consumer_cycle_state, true);

        // Update the Link TRB with the new cycle bit.
        let mut ltrb = LinkTrb::new();
        ltrb.set_trb_type(TrbType::Link);
        ltrb.set_ring_segment_pointer(0x100);
        ltrb.set_toggle_cycle(true);
        ltrb.set_cycle(true); // Producer cycle state is now 1.
        gm.write_obj_at_addr(ltrb, GuestAddress(0x100 + trb_size))
            .unwrap();

        // Overwrite the first TRB again with a new one (data = 3)
        // with the new producer cycle bit state (false).
        let mut trb = NormalTrb::new();
        trb.set_trb_type(TrbType::Normal);
        trb.set_data_buffer(3);
        trb.set_cycle(false); // Link TRB toggled the cycle.
        gm.write_obj_at_addr(trb.clone(), GuestAddress(0x100))
            .unwrap();

        // Read new transfer descriptor.
        let descriptor = transfer_ring
            .dequeue_transfer_descriptor()
            .unwrap()
            .unwrap();
        assert_eq!(descriptor.len(), 1);
        assert_eq!(descriptor[0].trb.get_parameter(), 3);

        assert_eq!(transfer_ring.consumer_cycle_state, false);
    }
}
<|endoftext|>