<commit_before>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

/// - The PubSub server sets up TCP listeners on 2 ports, the "subscriber" port and the "publisher"
///   port. Because both publishers and subscribers initiate their connections to the PubSub
///   server, the server requires no prior knowledge of either publishers or subscribers.
///
/// - Subscribers connect to the server on the server's "subscriber" port. Once a connection is
///   established, the server acts as the client of the Subscriber service, initially requesting
///   the topics the subscriber is interested in, and subsequently sending topical messages to the
///   subscriber.
///
/// - Publishers connect to the server on the "publisher" port and, once connected, they send
///   toopical messages via Publisher service to the server. The server then broadcasts each
///   messages to all clients subscribed to the topic of that message.
///
///       Subscriber                        Publisher                       PubSub Server
/// T1        |                                 |                                 |
/// T2        |-----Connect------------------------------------------------------>|
/// T3        |                                 |                                 |
/// T2        |<-------------------------------------------------------Topics-----|
/// T2        |-----(OK) Topics-------------------------------------------------->|
/// T3        |                                 |                                 |
/// T4        |                                 |-----Connect-------------------->|
/// T5        |                                 |                                 |
/// T6        |                                 |-----Publish-------------------->|
/// T7        |                                 |                                 |
/// T8        |<------------------------------------------------------Receive-----|
/// T9        |-----(OK) Receive------------------------------------------------->|
/// T10       |                                 |                                 |
/// T11       |                                 |<--------------(OK) Publish------|

use anyhow::anyhow;
use futures::{
    channel::oneshot,
    future::{self, AbortHandle},
    prelude::*,
};
use log::info;
use publisher::Publisher as _;
use std::{
    collections::HashMap,
    io,
    net::SocketAddr,
    sync::{Arc, Mutex, RwLock},
};
use subscriber::Subscriber as _;
use tarpc::{
    client, context,
    serde_transport::tcp,
    server::{self, Channel},
};
use tokio::net::ToSocketAddrs;
use tokio_serde::formats::Json;

pub mod subscriber {
    #[tarpc::service]
    pub trait Subscriber {
        async fn topics() -> Vec<String>;
        async fn receive(topic: String, message: String);
    }
}

pub mod publisher {
    #[tarpc::service]
    pub trait Publisher {
        async fn publish(topic: String, message: String);
    }
}

#[derive(Clone, Debug)]
struct Subscriber {
    local_addr: SocketAddr,
    topics: Vec<String>,
}

#[tarpc::server]
impl subscriber::Subscriber for Subscriber {
    async fn topics(self, _: context::Context) -> Vec<String> {
        self.topics.clone()
    }

    async fn receive(self, _: context::Context, topic: String, message: String) {
        info!(
            "[{}] received message on topic '{}': {}",
            self.local_addr, topic, message
        );
    }
}

struct SubscriberHandle(AbortHandle);

impl Drop for SubscriberHandle {
    fn drop(&mut self) {
        self.0.abort();
    }
}

impl Subscriber {
    async fn connect(
        publisher_addr: impl ToSocketAddrs,
        topics: Vec<String>,
    ) -> anyhow::Result<SubscriberHandle> {
        let publisher = tcp::connect(publisher_addr, Json::default()).await?;
        let local_addr = publisher.local_addr()?;
        let mut handler = server::BaseChannel::with_defaults(publisher)
            .respond_with(Subscriber { local_addr, topics }.serve());
        // The first request is for the topics being subscriibed to.
        match handler.next().await {
            Some(init_topics) => init_topics?.await,
            None => {
                return Err(anyhow!(
                    "[{}] Server never initialized the subscriber.",
                    local_addr
                ))
            }
        };
        let (handler, abort_handle) = future::abortable(handler.execute());
        tokio::spawn(async move {
            match handler.await {
                Ok(()) | Err(future::Aborted) => info!("[{}] subscriber shutdown.", local_addr),
            }
        });
        Ok(SubscriberHandle(abort_handle))
    }
}

#[derive(Debug)]
struct Subscription {
    subscriber: subscriber::SubscriberClient,
    topics: Vec<String>,
}

#[derive(Clone, Debug)]
struct Publisher {
    clients: Arc<Mutex<HashMap<SocketAddr, Subscription>>>,
    subscriptions: Arc<RwLock<HashMap<String, HashMap<SocketAddr, subscriber::SubscriberClient>>>>,
}

struct PublisherAddrs {
    publisher: SocketAddr,
    subscriptions: SocketAddr,
}

impl Publisher {
    async fn start(self) -> io::Result<PublisherAddrs> {
        let mut connecting_publishers = tcp::listen("localhost:0", Json::default).await?;

        let publisher_addrs = PublisherAddrs {
            publisher: connecting_publishers.local_addr(),
            subscriptions: self.clone().start_subscription_manager().await?,
        };

        info!("[{}] listening for publishers.", publisher_addrs.publisher);
        tokio::spawn(async move {
            // Because this is just an example, we know there will only be one publisher. In more
            // realistic code, this would be a loop to continually accept new publisher
            // connections.
            let publisher = connecting_publishers.next().await.unwrap().unwrap();
            info!("[{}] publisher connected.", publisher.peer_addr().unwrap());

            server::BaseChannel::with_defaults(publisher)
                .respond_with(self.serve())
                .execute()
                .await
        });

        Ok(publisher_addrs)
    }

    async fn start_subscription_manager(self) -> io::Result<SocketAddr> {
        let mut connecting_subscribers = tcp::listen("localhost:0", Json::default)
            .await?
            .filter_map(|r| future::ready(r.ok()));
        let new_subscriber_addr = connecting_subscribers.get_ref().local_addr();
        info!("[{}] listening for subscribers.", new_subscriber_addr);

        tokio::spawn(async move {
            while let Some(conn) = connecting_subscribers.next().await {
                let subscriber_addr = conn.peer_addr().unwrap();

                let tarpc::client::NewClient {
                    client: mut subscriber,
                    dispatch,
                } = subscriber::SubscriberClient::new(client::Config::default(), conn);
                let (ready_tx, ready) = oneshot::channel();
                self.clone()
                    .start_subscriber_gc(subscriber_addr, dispatch, ready);

                // Populate the topics
                if let Ok(topics) = subscriber.topics(context::current()).await {
                    self.clients.lock().unwrap().insert(
                        subscriber_addr,
                        Subscription {
                            subscriber: subscriber.clone(),
                            topics: topics.clone(),
                        },
                    );

                    info!("[{}] subscribed to topics: {:?}", subscriber_addr, topics);
                    let mut subscriptions = self.subscriptions.write().unwrap();
                    for topic in topics {
                        subscriptions
                            .entry(topic)
                            .or_insert_with(HashMap::new)
                            .insert(subscriber_addr, subscriber.clone());
                    }
                }
                // Signal that initialization is done.
                ready_tx.send(()).unwrap();
            }
        });

        Ok(new_subscriber_addr)
    }

    fn start_subscriber_gc(
        self,
        subscriber_addr: SocketAddr,
        client_dispatch: impl Future<Output = anyhow::Result<()>> + Send + 'static,
        subscriber_ready: oneshot::Receiver<()>,
    ) {
        tokio::spawn(async move {
            if let Err(e) = client_dispatch.await {
                info!("[{}] subscriber connection broken: {:?}", subscriber_addr, e)
            }
            // Don't clean up the subscriber until initialization is done.
            let _ = subscriber_ready.await;
            if let Some(subscription) = self.clients.lock().unwrap().remove(&subscriber_addr) {
                info!("[{} unsubscribing from topics: {:?}", subscriber_addr, subscription.topics);
                let mut subscriptions = self.subscriptions.write().unwrap();
                for topic in subscription.topics {
                    let subscribers = subscriptions.get_mut(&topic).unwrap();
                    subscribers.remove(&subscriber_addr);
                    if subscribers.is_empty() {
                        subscriptions.remove(&topic);
                    }
                }
            }
        });
    }
}

#[tarpc::server]
impl publisher::Publisher for Publisher {
    async fn publish(self, _: context::Context, topic: String, message: String) {
        info!("received message to publish.");
        let mut subscribers = match self.subscriptions.read().unwrap().get(&topic) {
            None => return,
            Some(subscriptions) => subscriptions.clone(),
        };
        let mut publications = Vec::new();
        for client in subscribers.values_mut() {
            publications.push(client.receive(context::current(), topic.clone(), message.clone()));
        }
        // Ignore failing subscribers. In a real pubsub, you'd want to continually retry until
        // subscribers ack. Of course, a lot would be different in a real pubsub :)
        for response in future::join_all(publications).await {
            if let Err(e) = response {
                info!("failed to broadcast to subscriber: {}", e);
            }
        }
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::init();

    let clients = Arc::new(Mutex::new(HashMap::new()));
    let addrs = Publisher {
        clients,
        subscriptions: Arc::new(RwLock::new(HashMap::new())),
    }
    .start()
    .await?;

    let _subscriber0 = Subscriber::connect(
        addrs.subscriptions,
        vec!["calculus".into(), "cool shorts".into()],
    )
    .await?;

    let _subscriber1 = Subscriber::connect(
        addrs.subscriptions,
        vec!["cool shorts".into(), "history".into()],
    )
    .await?;

    let mut publisher = publisher::PublisherClient::new(
        client::Config::default(),
        tcp::connect(addrs.publisher, Json::default()).await?,
    )
    .spawn()?;

    publisher
        .publish(context::current(), "calculus".into(), "sqrt(2)".into())
        .await?;

    publisher
        .publish(
            context::current(),
            "cool shorts".into(),
            "hello to all".into(),
        )
        .await?;

    publisher
        .publish(context::current(), "history".into(), "napoleon".to_string())
        .await?;

    drop(_subscriber0);

    publisher
        .publish(
            context::current(),
            "cool shorts".into(),
            "hello to who?".into(),
        )
        .await?;

    info!("done.");

    Ok(())
}
<commit_msg>Minor refactor<commit_after>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

/// - The PubSub server sets up TCP listeners on 2 ports, the "subscriber" port and the "publisher"
///   port. Because both publishers and subscribers initiate their connections to the PubSub
///   server, the server requires no prior knowledge of either publishers or subscribers.
///
/// - Subscribers connect to the server on the server's "subscriber" port. Once a connection is
///   established, the server acts as the client of the Subscriber service, initially requesting
///   the topics the subscriber is interested in, and subsequently sending topical messages to the
///   subscriber.
///
/// - Publishers connect to the server on the "publisher" port and, once connected, they send
///   topical messages via Publisher service to the server. The server then broadcasts each
///   messages to all clients subscribed to the topic of that message.
///
///       Subscriber                        Publisher                       PubSub Server
/// T1        |                                 |                                 |
/// T2        |-----Connect------------------------------------------------------>|
/// T3        |                                 |                                 |
/// T2        |<-------------------------------------------------------Topics-----|
/// T2        |-----(OK) Topics-------------------------------------------------->|
/// T3        |                                 |                                 |
/// T4        |                                 |-----Connect-------------------->|
/// T5        |                                 |                                 |
/// T6        |                                 |-----Publish-------------------->|
/// T7        |                                 |                                 |
/// T8        |<------------------------------------------------------Receive-----|
/// T9        |-----(OK) Receive------------------------------------------------->|
/// T10       |                                 |                                 |
/// T11       |                                 |<--------------(OK) Publish------|
use anyhow::anyhow;
use futures::{
    channel::oneshot,
    future::{self, AbortHandle},
    prelude::*,
};
use log::info;
use publisher::Publisher as _;
use std::{
    collections::HashMap,
    io,
    net::SocketAddr,
    sync::{Arc, Mutex, RwLock},
};
use subscriber::Subscriber as _;
use tarpc::{
    client, context,
    serde_transport::tcp,
    server::{self, Channel},
};
use tokio::net::ToSocketAddrs;
use tokio_serde::formats::Json;

pub mod subscriber {
    #[tarpc::service]
    pub trait Subscriber {
        async fn topics() -> Vec<String>;
        async fn receive(topic: String, message: String);
    }
}

pub mod publisher {
    #[tarpc::service]
    pub trait Publisher {
        async fn publish(topic: String, message: String);
    }
}

#[derive(Clone, Debug)]
struct Subscriber {
    local_addr: SocketAddr,
    topics: Vec<String>,
}

#[tarpc::server]
impl subscriber::Subscriber for Subscriber {
    async fn topics(self, _: context::Context) -> Vec<String> {
        self.topics.clone()
    }

    async fn receive(self, _: context::Context, topic: String, message: String) {
        info!(
            "[{}] received message on topic '{}': {}",
            self.local_addr, topic, message
        );
    }
}

struct SubscriberHandle(AbortHandle);

impl Drop for SubscriberHandle {
    fn drop(&mut self) {
        self.0.abort();
    }
}

impl Subscriber {
    async fn connect(
        publisher_addr: impl ToSocketAddrs,
        topics: Vec<String>,
    ) -> anyhow::Result<SubscriberHandle> {
        let publisher = tcp::connect(publisher_addr, Json::default()).await?;
        let local_addr = publisher.local_addr()?;
        let mut handler = server::BaseChannel::with_defaults(publisher)
            .respond_with(Subscriber { local_addr, topics }.serve());
        // The first request is for the topics being subscriibed to.
        match handler.next().await {
            Some(init_topics) => init_topics?.await,
            None => {
                return Err(anyhow!(
                    "[{}] Server never initialized the subscriber.",
                    local_addr
                ))
            }
        };
        let (handler, abort_handle) = future::abortable(handler.execute());
        tokio::spawn(async move {
            match handler.await {
                Ok(()) | Err(future::Aborted) => info!("[{}] subscriber shutdown.", local_addr),
            }
        });
        Ok(SubscriberHandle(abort_handle))
    }
}

#[derive(Debug)]
struct Subscription {
    subscriber: subscriber::SubscriberClient,
    topics: Vec<String>,
}

#[derive(Clone, Debug)]
struct Publisher {
    clients: Arc<Mutex<HashMap<SocketAddr, Subscription>>>,
    subscriptions: Arc<RwLock<HashMap<String, HashMap<SocketAddr, subscriber::SubscriberClient>>>>,
}

struct PublisherAddrs {
    publisher: SocketAddr,
    subscriptions: SocketAddr,
}

impl Publisher {
    async fn start(self) -> io::Result<PublisherAddrs> {
        let mut connecting_publishers = tcp::listen("localhost:0", Json::default).await?;

        let publisher_addrs = PublisherAddrs {
            publisher: connecting_publishers.local_addr(),
            subscriptions: self.clone().start_subscription_manager().await?,
        };

        info!("[{}] listening for publishers.", publisher_addrs.publisher);
        tokio::spawn(async move {
            // Because this is just an example, we know there will only be one publisher. In more
            // realistic code, this would be a loop to continually accept new publisher
            // connections.
            let publisher = connecting_publishers.next().await.unwrap().unwrap();
            info!("[{}] publisher connected.", publisher.peer_addr().unwrap());

            server::BaseChannel::with_defaults(publisher)
                .respond_with(self.serve())
                .execute()
                .await
        });

        Ok(publisher_addrs)
    }

    async fn start_subscription_manager(mut self) -> io::Result<SocketAddr> {
        let mut connecting_subscribers = tcp::listen("localhost:0", Json::default)
            .await?
            .filter_map(|r| future::ready(r.ok()));
        let new_subscriber_addr = connecting_subscribers.get_ref().local_addr();
        info!("[{}] listening for subscribers.", new_subscriber_addr);

        tokio::spawn(async move {
            while let Some(conn) = connecting_subscribers.next().await {
                let subscriber_addr = conn.peer_addr().unwrap();

                let tarpc::client::NewClient {
                    client: subscriber,
                    dispatch,
                } = subscriber::SubscriberClient::new(client::Config::default(), conn);
                let (ready_tx, ready) = oneshot::channel();
                self.clone()
                    .start_subscriber_gc(subscriber_addr, dispatch, ready);

                // Populate the topics
                self.initialize_subscription(subscriber_addr, subscriber)
                    .await;

                // Signal that initialization is done.
                ready_tx.send(()).unwrap();
            }
        });

        Ok(new_subscriber_addr)
    }

    async fn initialize_subscription(
        &mut self,
        subscriber_addr: SocketAddr,
        mut subscriber: subscriber::SubscriberClient,
    ) {
        // Populate the topics
        if let Ok(topics) = subscriber.topics(context::current()).await {
            self.clients.lock().unwrap().insert(
                subscriber_addr,
                Subscription {
                    subscriber: subscriber.clone(),
                    topics: topics.clone(),
                },
            );

            info!("[{}] subscribed to topics: {:?}", subscriber_addr, topics);
            let mut subscriptions = self.subscriptions.write().unwrap();
            for topic in topics {
                subscriptions
                    .entry(topic)
                    .or_insert_with(HashMap::new)
                    .insert(subscriber_addr, subscriber.clone());
            }
        }
    }

    fn start_subscriber_gc(
        self,
        subscriber_addr: SocketAddr,
        client_dispatch: impl Future<Output = anyhow::Result<()>> + Send + 'static,
        subscriber_ready: oneshot::Receiver<()>,
    ) {
        tokio::spawn(async move {
            if let Err(e) = client_dispatch.await {
                info!(
                    "[{}] subscriber connection broken: {:?}",
                    subscriber_addr, e
                )
            }
            // Don't clean up the subscriber until initialization is done.
            let _ = subscriber_ready.await;
            if let Some(subscription) = self.clients.lock().unwrap().remove(&subscriber_addr) {
                info!(
                    "[{} unsubscribing from topics: {:?}",
                    subscriber_addr, subscription.topics
                );
                let mut subscriptions = self.subscriptions.write().unwrap();
                for topic in subscription.topics {
                    let subscribers = subscriptions.get_mut(&topic).unwrap();
                    subscribers.remove(&subscriber_addr);
                    if subscribers.is_empty() {
                        subscriptions.remove(&topic);
                    }
                }
            }
        });
    }
}

#[tarpc::server]
impl publisher::Publisher for Publisher {
    async fn publish(self, _: context::Context, topic: String, message: String) {
        info!("received message to publish.");
        let mut subscribers = match self.subscriptions.read().unwrap().get(&topic) {
            None => return,
            Some(subscriptions) => subscriptions.clone(),
        };
        let mut publications = Vec::new();
        for client in subscribers.values_mut() {
            publications.push(client.receive(context::current(), topic.clone(), message.clone()));
        }
        // Ignore failing subscribers. In a real pubsub, you'd want to continually retry until
        // subscribers ack. Of course, a lot would be different in a real pubsub :)
        for response in future::join_all(publications).await {
            if let Err(e) = response {
                info!("failed to broadcast to subscriber: {}", e);
            }
        }
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::init();

    let clients = Arc::new(Mutex::new(HashMap::new()));
    let addrs = Publisher {
        clients,
        subscriptions: Arc::new(RwLock::new(HashMap::new())),
    }
    .start()
    .await?;

    let _subscriber0 = Subscriber::connect(
        addrs.subscriptions,
        vec!["calculus".into(), "cool shorts".into()],
    )
    .await?;

    let _subscriber1 = Subscriber::connect(
        addrs.subscriptions,
        vec!["cool shorts".into(), "history".into()],
    )
    .await?;

    let mut publisher = publisher::PublisherClient::new(
        client::Config::default(),
        tcp::connect(addrs.publisher, Json::default()).await?,
    )
    .spawn()?;

    publisher
        .publish(context::current(), "calculus".into(), "sqrt(2)".into())
        .await?;

    publisher
        .publish(
            context::current(),
            "cool shorts".into(),
            "hello to all".into(),
        )
        .await?;

    publisher
        .publish(context::current(), "history".into(), "napoleon".to_string())
        .await?;

    drop(_subscriber0);

    publisher
        .publish(
            context::current(),
            "cool shorts".into(),
            "hello to who?".into(),
        )
        .await?;

    info!("done.");

    Ok(())
}
<|endoftext|>


<commit_before>// Repeat the function f, n times.
fn repeat<F>(f: &F, n: u32)
    where F: Fn() {
    for _ in 0..n {
        f();
    }
}

fn static_fn() {
    print!("Static ");
}

fn main() {

    // Repeat a static function.
    repeat(&static_fn, 4);

    println!("");

    // Repeat an anonymous closure.
    repeat(&|| print!("Closure "), 5);
}
<commit_msg>Added mutable closure test<commit_after>// Repeat the function f, n times.
fn repeat<F>(f: &F, n: u32)
    where F: Fn() {
    for _ in 0..n {
        f();
    }
}

// Repeat the mutable function f, n times.
fn repeat_mut<F>(f: &mut F, n: u32)
    where F: FnMut() {
    for _ in 0..n {
        f();
    }
}

fn static_fn() {
    print!("Static ");
}

fn main() {

    // Repeat a static function.
    repeat(&static_fn, 4);

    println!("");

    // Repeat an anonymous closure.
    repeat(&|| print!("Closure "), 5);

    println!("");

    // Repeat a mutable closure (can modify local variables).
    let mut x = 1;
    println!("X is {}", x);
    repeat_mut(&mut || x = x + 1, 5);
    println!("X is now {}", x);
}

#[test]
fn test_closure() {
    let mut x = 1;

    repeat_mut(&mut || x = x + 1, 5);

    assert_eq!(x, 6);
}<|endoftext|>


<commit_before><commit_msg>Implement Index for Packed<commit_after><|endoftext|>


<commit_before>use syntax_definition::*;
use scope::*;
use yaml_load::*;
use theme::theme::{Theme, ParseThemeError};
use theme::settings::*;

use std::path::{Path, PathBuf};
use std::io::{Error as IoError, BufReader};
use walkdir::WalkDir;
use std::io::{self, Read};
use std::fs::File;
use walkdir;
use std::ops::DerefMut;
use std::mem;
use std::rc::Rc;
use std::ascii::AsciiExt;

#[derive(Debug)]
pub struct PackageSet {
    pub syntaxes: Vec<SyntaxDefinition>,
}

#[derive(Debug)]
pub enum PackageError {
    WalkDir(walkdir::Error),
    Io(io::Error),
    ParseSyntax(ParseSyntaxError),
    ParseTheme(ParseThemeError),
    ReadSettings(SettingsError),
}

impl From<SettingsError> for PackageError {
    fn from(error: SettingsError) -> PackageError {
        PackageError::ReadSettings(error)
    }
}

impl From<IoError> for PackageError {
    fn from(error: IoError) -> PackageError {
        PackageError::Io(error)
    }
}

impl From<ParseThemeError> for PackageError {
    fn from(error: ParseThemeError) -> PackageError {
        PackageError::ParseTheme(error)
    }
}

impl From<ParseSyntaxError> for PackageError {
    fn from(error: ParseSyntaxError) -> PackageError {
        PackageError::ParseSyntax(error)
    }
}

fn load_syntax_file(p: &Path, lines_include_newline: bool) -> Result<SyntaxDefinition, PackageError> {
    let mut f = try!(File::open(p));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));

    Ok(try!(SyntaxDefinition::load_from_str(&s, lines_include_newline)))
}

impl PackageSet {
    pub fn new() -> PackageSet {
        PackageSet { syntaxes: Vec::new() }
    }

    /// Convenience constructor calling `new` and then `load_syntaxes` on the resulting set
    /// defaults to lines given not including newline characters, see the
    /// `load_syntaxes` method docs for an explanation as to why this might not be the best.
    pub fn load_from_folder<P: AsRef<Path>>(folder: P) -> Result<PackageSet, PackageError> {
        let mut ps = Self::new();
        try!(ps.load_syntaxes(folder, false));
        Ok(ps)
    }

    /// Returns all the themes found in a folder, good for enumerating before loading one with get_theme
    pub fn discover_themes<P: AsRef<Path>>(folder: P) -> Result<Vec<PathBuf>, PackageError> {
        let mut themes = Vec::new();
        for entry in WalkDir::new(folder) {
            let entry = try!(entry.map_err(|e| PackageError::WalkDir(e)));
            if entry.path().extension().map(|e| e == "tmTheme").unwrap_or(false) {
                themes.push(entry.path().to_owned());
            }
        }
        Ok(themes)
    }

    /// Loads all the .sublime-syntax files in a folder and links them together into this package set
    ///
    /// The `lines_include_newline` parameter is used to work around the fact that Sublime Text normally
    /// passes line strings including newline characters (`\n`) to its regex engine. This results in many
    /// syntaxes having regexes matching `\n`, which doesn't work if you don't pass in newlines.
    /// It is recommended that if you can you pass in lines with newlines if you can and pass `true` for this parameter.
    /// If that is inconvenient pass `false` and the loader will do some hacky find and replaces on the
    /// match regexes that seem to work for the default syntax set, but may not work for any other syntaxes.
    ///
    /// In the future I might include a "slow mode" that copies the lines passed in and appends a newline if there isn't one.
    /// but in the interest of performance currently this hacky fix will have to do.
    pub fn load_syntaxes<P: AsRef<Path>>(&mut self, folder: P, lines_include_newline: bool) -> Result<(), PackageError> {
        for entry in WalkDir::new(folder) {
            let entry = try!(entry.map_err(|e| PackageError::WalkDir(e)));
            if entry.path().extension().map(|e| e == "sublime-syntax").unwrap_or(false) {
                // println!("{}", entry.path().display());
                self.syntaxes.push(try!(load_syntax_file(entry.path(), lines_include_newline)));
            }
        }
        self.link_syntaxes();
        Ok(())
    }

    pub fn find_syntax_by_scope<'a>(&'a self, scope: Scope) -> Option<&'a SyntaxDefinition> {
        self.syntaxes.iter().find(|&s| s.scope == scope)
    }

    pub fn find_syntax_by_name<'a>(&'a self, name: &str) -> Option<&'a SyntaxDefinition> {
        self.syntaxes.iter().find(|&s| name == &s.name)
    }

    pub fn find_syntax_by_extension<'a>(&'a self, extension: &str) -> Option<&'a SyntaxDefinition> {
        self.syntaxes.iter().find(|&s| s.file_extensions.iter().any(|e| e == extension))
    }

    /// Searches for a syntax first by extension and then by case-insensitive name
    /// useful for things like Github-flavoured-markdown code block highlighting where
    /// all you have to go on is a short token given by the user
    pub fn find_syntax_by_token<'a>(&'a self, s: &str) -> Option<&'a SyntaxDefinition> {
        {
            let ext_res = self.find_syntax_by_extension(s);
            if ext_res.is_some() {
                return ext_res;
            }
        }
        let lower = s.to_ascii_lowercase();
        self.syntaxes.iter().find(|&s| lower == s.name.to_ascii_lowercase())
    }

    fn link_syntaxes(&mut self) {
        for syntax in self.syntaxes.iter() {
            for (_, ref context_ptr) in syntax.contexts.iter() {
                let mut mut_ref = context_ptr.borrow_mut();
                self.link_context(syntax, mut_ref.deref_mut());
            }
        }
    }

    fn link_context(&self, syntax: &SyntaxDefinition, context: &mut Context) {
        for pattern in context.patterns.iter_mut() {
            match *pattern {
                Pattern::Match(ref mut match_pat) => self.link_match_pat(syntax, match_pat),
                Pattern::Include(ref mut context_ref) => self.link_ref(syntax, context_ref),
            }
        }
    }

    fn link_ref(&self, syntax: &SyntaxDefinition, context_ref: &mut ContextReference) {
        // println!("{:?}", context_ref);
        use syntax_definition::ContextReference::*;
        let maybe_new_context = match *context_ref {
            Named(ref s) => syntax.contexts.get(s),
            Inline(ref context_ptr) => {
                let mut mut_ref = context_ptr.borrow_mut();
                self.link_context(syntax, mut_ref.deref_mut());
                None
            }
            ByScope { scope, ref sub_context } => {
                let other_syntax = self.find_syntax_by_scope(scope);
                let context_name = sub_context.as_ref().map(|x| &**x).unwrap_or("main");
                other_syntax.and_then(|s| s.contexts.get(context_name))
            }
            File { ref name, ref sub_context } => {
                let other_syntax = self.find_syntax_by_name(name);
                let context_name = sub_context.as_ref().map(|x| &**x).unwrap_or("main");
                other_syntax.and_then(|s| s.contexts.get(context_name))
            }
            Direct(_) => None,
        };
        if let Some(new_context) = maybe_new_context {
            let mut new_ref = Direct(Rc::downgrade(&new_context));
            mem::swap(context_ref, &mut new_ref);
        }
    }

    fn link_match_pat(&self, syntax: &SyntaxDefinition, match_pat: &mut MatchPattern) {
        let maybe_context_refs = match match_pat.operation {
            MatchOperation::Push(ref mut context_refs) => Some(context_refs),
            MatchOperation::Set(ref mut context_refs) => Some(context_refs),
            MatchOperation::Pop | MatchOperation::None => None,
        };
        if let Some(context_refs) = maybe_context_refs {
            for context_ref in context_refs.iter_mut() {
                self.link_ref(syntax, context_ref);
            }
        }
        if let Some(ref context_ptr) = match_pat.with_prototype {
            let mut mut_ref = context_ptr.borrow_mut();
            self.link_context(syntax, mut_ref.deref_mut());
        }
    }

    fn read_file(path: &Path) -> Result<BufReader<File>, PackageError> {
        let reader = try!(File::open(path));
        Ok(BufReader::new(reader))
    }

    fn read_plist(path: &Path) -> Result<Settings, PackageError> {
        Ok(try!(read_plist(try!(Self::read_file(path)))))
    }

    pub fn get_theme<P: AsRef<Path>>(path: P) -> Result<Theme, PackageError> {
        Ok(try!(Theme::parse_settings(try!(Self::read_plist(path.as_ref())))))
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn can_load() {
        use package_set::PackageSet;
        use syntax_definition::*;
        use scope::*;
        let ps = PackageSet::load_from_folder("testdata/Packages").unwrap();
        let rails_scope = Scope::new("source.ruby.rails").unwrap();
        let syntax = ps.find_syntax_by_name("Ruby on Rails").unwrap();
        assert_eq!(&ps.find_syntax_by_extension("rake").unwrap().name, "Ruby");
        assert_eq!(&ps.find_syntax_by_token("ruby").unwrap().name, "Ruby");
        // println!("{:#?}", syntax);
        assert_eq!(syntax.scope, rails_scope);
        // assert!(false);
        let main_context = syntax.contexts.get("main").unwrap();
        let count = context_iter(main_context.clone()).count();
        assert_eq!(count, 91);
    }
    #[test]
    fn can_parse_common_themes() {
        use package_set::PackageSet;
        use theme::style::Color;
        let theme_paths = PackageSet::discover_themes("testdata/themes.tmbundle").unwrap();
        for theme_path in theme_paths.iter() {
            println!("{:?}", theme_path);
            PackageSet::get_theme(theme_path).unwrap();
        }

        let theme = PackageSet::get_theme("testdata/themes.tmbundle/Themes/Amy.tmTheme").unwrap();
        assert_eq!(theme.name.unwrap(), "Amy");
        assert_eq!(theme.settings.selection.unwrap(),
                   Color {
                       r: 0x80,
                       g: 0x00,
                       b: 0x00,
                       a: 0x80,
                   });
        assert_eq!(theme.scopes[0].style.foreground.unwrap(),
                   Color {
                       r: 0x40,
                       g: 0x40,
                       b: 0x80,
                       a: 0xFF,
                   });
        // assert!(false);
    }
}
<commit_msg>Cargo fmt<commit_after>use syntax_definition::*;
use scope::*;
use yaml_load::*;
use theme::theme::{Theme, ParseThemeError};
use theme::settings::*;

use std::path::{Path, PathBuf};
use std::io::{Error as IoError, BufReader};
use walkdir::WalkDir;
use std::io::{self, Read};
use std::fs::File;
use walkdir;
use std::ops::DerefMut;
use std::mem;
use std::rc::Rc;
use std::ascii::AsciiExt;

#[derive(Debug)]
pub struct PackageSet {
    pub syntaxes: Vec<SyntaxDefinition>,
}

#[derive(Debug)]
pub enum PackageError {
    WalkDir(walkdir::Error),
    Io(io::Error),
    ParseSyntax(ParseSyntaxError),
    ParseTheme(ParseThemeError),
    ReadSettings(SettingsError),
}

impl From<SettingsError> for PackageError {
    fn from(error: SettingsError) -> PackageError {
        PackageError::ReadSettings(error)
    }
}

impl From<IoError> for PackageError {
    fn from(error: IoError) -> PackageError {
        PackageError::Io(error)
    }
}

impl From<ParseThemeError> for PackageError {
    fn from(error: ParseThemeError) -> PackageError {
        PackageError::ParseTheme(error)
    }
}

impl From<ParseSyntaxError> for PackageError {
    fn from(error: ParseSyntaxError) -> PackageError {
        PackageError::ParseSyntax(error)
    }
}

fn load_syntax_file(p: &Path,
                    lines_include_newline: bool)
                    -> Result<SyntaxDefinition, PackageError> {
    let mut f = try!(File::open(p));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));

    Ok(try!(SyntaxDefinition::load_from_str(&s, lines_include_newline)))
}

impl PackageSet {
    pub fn new() -> PackageSet {
        PackageSet { syntaxes: Vec::new() }
    }

    /// Convenience constructor calling `new` and then `load_syntaxes` on the resulting set
    /// defaults to lines given not including newline characters, see the
    /// `load_syntaxes` method docs for an explanation as to why this might not be the best.
    pub fn load_from_folder<P: AsRef<Path>>(folder: P) -> Result<PackageSet, PackageError> {
        let mut ps = Self::new();
        try!(ps.load_syntaxes(folder, false));
        Ok(ps)
    }

    /// Returns all the themes found in a folder, good for enumerating before loading one with get_theme
    pub fn discover_themes<P: AsRef<Path>>(folder: P) -> Result<Vec<PathBuf>, PackageError> {
        let mut themes = Vec::new();
        for entry in WalkDir::new(folder) {
            let entry = try!(entry.map_err(|e| PackageError::WalkDir(e)));
            if entry.path().extension().map(|e| e == "tmTheme").unwrap_or(false) {
                themes.push(entry.path().to_owned());
            }
        }
        Ok(themes)
    }

    /// Loads all the .sublime-syntax files in a folder and links them together into this package set
    ///
    /// The `lines_include_newline` parameter is used to work around the fact that Sublime Text normally
    /// passes line strings including newline characters (`\n`) to its regex engine. This results in many
    /// syntaxes having regexes matching `\n`, which doesn't work if you don't pass in newlines.
    /// It is recommended that if you can you pass in lines with newlines if you can and pass `true` for this parameter.
    /// If that is inconvenient pass `false` and the loader will do some hacky find and replaces on the
    /// match regexes that seem to work for the default syntax set, but may not work for any other syntaxes.
    ///
    /// In the future I might include a "slow mode" that copies the lines passed in and appends a newline if there isn't one.
    /// but in the interest of performance currently this hacky fix will have to do.
    pub fn load_syntaxes<P: AsRef<Path>>(&mut self,
                                         folder: P,
                                         lines_include_newline: bool)
                                         -> Result<(), PackageError> {
        for entry in WalkDir::new(folder) {
            let entry = try!(entry.map_err(|e| PackageError::WalkDir(e)));
            if entry.path().extension().map(|e| e == "sublime-syntax").unwrap_or(false) {
                // println!("{}", entry.path().display());
                self.syntaxes.push(try!(load_syntax_file(entry.path(), lines_include_newline)));
            }
        }
        self.link_syntaxes();
        Ok(())
    }

    pub fn find_syntax_by_scope<'a>(&'a self, scope: Scope) -> Option<&'a SyntaxDefinition> {
        self.syntaxes.iter().find(|&s| s.scope == scope)
    }

    pub fn find_syntax_by_name<'a>(&'a self, name: &str) -> Option<&'a SyntaxDefinition> {
        self.syntaxes.iter().find(|&s| name == &s.name)
    }

    pub fn find_syntax_by_extension<'a>(&'a self, extension: &str) -> Option<&'a SyntaxDefinition> {
        self.syntaxes.iter().find(|&s| s.file_extensions.iter().any(|e| e == extension))
    }

    /// Searches for a syntax first by extension and then by case-insensitive name
    /// useful for things like Github-flavoured-markdown code block highlighting where
    /// all you have to go on is a short token given by the user
    pub fn find_syntax_by_token<'a>(&'a self, s: &str) -> Option<&'a SyntaxDefinition> {
        {
            let ext_res = self.find_syntax_by_extension(s);
            if ext_res.is_some() {
                return ext_res;
            }
        }
        let lower = s.to_ascii_lowercase();
        self.syntaxes.iter().find(|&s| lower == s.name.to_ascii_lowercase())
    }

    fn link_syntaxes(&mut self) {
        for syntax in self.syntaxes.iter() {
            for (_, ref context_ptr) in syntax.contexts.iter() {
                let mut mut_ref = context_ptr.borrow_mut();
                self.link_context(syntax, mut_ref.deref_mut());
            }
        }
    }

    fn link_context(&self, syntax: &SyntaxDefinition, context: &mut Context) {
        for pattern in context.patterns.iter_mut() {
            match *pattern {
                Pattern::Match(ref mut match_pat) => self.link_match_pat(syntax, match_pat),
                Pattern::Include(ref mut context_ref) => self.link_ref(syntax, context_ref),
            }
        }
    }

    fn link_ref(&self, syntax: &SyntaxDefinition, context_ref: &mut ContextReference) {
        // println!("{:?}", context_ref);
        use syntax_definition::ContextReference::*;
        let maybe_new_context = match *context_ref {
            Named(ref s) => syntax.contexts.get(s),
            Inline(ref context_ptr) => {
                let mut mut_ref = context_ptr.borrow_mut();
                self.link_context(syntax, mut_ref.deref_mut());
                None
            }
            ByScope { scope, ref sub_context } => {
                let other_syntax = self.find_syntax_by_scope(scope);
                let context_name = sub_context.as_ref().map(|x| &**x).unwrap_or("main");
                other_syntax.and_then(|s| s.contexts.get(context_name))
            }
            File { ref name, ref sub_context } => {
                let other_syntax = self.find_syntax_by_name(name);
                let context_name = sub_context.as_ref().map(|x| &**x).unwrap_or("main");
                other_syntax.and_then(|s| s.contexts.get(context_name))
            }
            Direct(_) => None,
        };
        if let Some(new_context) = maybe_new_context {
            let mut new_ref = Direct(Rc::downgrade(&new_context));
            mem::swap(context_ref, &mut new_ref);
        }
    }

    fn link_match_pat(&self, syntax: &SyntaxDefinition, match_pat: &mut MatchPattern) {
        let maybe_context_refs = match match_pat.operation {
            MatchOperation::Push(ref mut context_refs) => Some(context_refs),
            MatchOperation::Set(ref mut context_refs) => Some(context_refs),
            MatchOperation::Pop | MatchOperation::None => None,
        };
        if let Some(context_refs) = maybe_context_refs {
            for context_ref in context_refs.iter_mut() {
                self.link_ref(syntax, context_ref);
            }
        }
        if let Some(ref context_ptr) = match_pat.with_prototype {
            let mut mut_ref = context_ptr.borrow_mut();
            self.link_context(syntax, mut_ref.deref_mut());
        }
    }

    fn read_file(path: &Path) -> Result<BufReader<File>, PackageError> {
        let reader = try!(File::open(path));
        Ok(BufReader::new(reader))
    }

    fn read_plist(path: &Path) -> Result<Settings, PackageError> {
        Ok(try!(read_plist(try!(Self::read_file(path)))))
    }

    pub fn get_theme<P: AsRef<Path>>(path: P) -> Result<Theme, PackageError> {
        Ok(try!(Theme::parse_settings(try!(Self::read_plist(path.as_ref())))))
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn can_load() {
        use package_set::PackageSet;
        use syntax_definition::*;
        use scope::*;
        let ps = PackageSet::load_from_folder("testdata/Packages").unwrap();
        let rails_scope = Scope::new("source.ruby.rails").unwrap();
        let syntax = ps.find_syntax_by_name("Ruby on Rails").unwrap();
        assert_eq!(&ps.find_syntax_by_extension("rake").unwrap().name, "Ruby");
        assert_eq!(&ps.find_syntax_by_token("ruby").unwrap().name, "Ruby");
        // println!("{:#?}", syntax);
        assert_eq!(syntax.scope, rails_scope);
        // assert!(false);
        let main_context = syntax.contexts.get("main").unwrap();
        let count = context_iter(main_context.clone()).count();
        assert_eq!(count, 91);
    }
    #[test]
    fn can_parse_common_themes() {
        use package_set::PackageSet;
        use theme::style::Color;
        let theme_paths = PackageSet::discover_themes("testdata/themes.tmbundle").unwrap();
        for theme_path in theme_paths.iter() {
            println!("{:?}", theme_path);
            PackageSet::get_theme(theme_path).unwrap();
        }

        let theme = PackageSet::get_theme("testdata/themes.tmbundle/Themes/Amy.tmTheme").unwrap();
        assert_eq!(theme.name.unwrap(), "Amy");
        assert_eq!(theme.settings.selection.unwrap(),
                   Color {
                       r: 0x80,
                       g: 0x00,
                       b: 0x00,
                       a: 0x80,
                   });
        assert_eq!(theme.scopes[0].style.foreground.unwrap(),
                   Color {
                       r: 0x40,
                       g: 0x40,
                       b: 0x80,
                       a: 0xFF,
                   });
        // assert!(false);
    }
}
<|endoftext|>


<commit_before>use action::Action;
use context::base::BaseContext;
use dispatcher::Response;
use dispatcher::response::ResponseSender;
use message::{Message, MessageBuilder};

use handlebars::Context;
use std::borrow::Borrow;
use std::collections::BTreeMap;
use state::State;
use self::error::Error;
use self::renderer_context::RendererContext;

mod error;
mod renderer_context;
#[cfg(test)]
mod test;

pub const CONTEXT_UUID: &'static str = "context_uuid";
pub const CONTEXT_NAME: &'static str = "context_name";
pub const CONTEXT_LEN: &'static str = "context_len";
pub const MESSAGES: &'static str = "messages";
const MESSAGE: &'static str = "MESSAGE";

use handlebars::Template;
use handlebars::Handlebars;

use super::ActionType;
use super::ExecCondition;

mod deser;
mod builder;

pub use self::builder::MessageActionBuilder;

pub struct MessageAction {
    pub uuid: String,
    pub name: Option<String>,
    pub message: Template,
    pub values: Handlebars,
    pub when: ExecCondition,
    pub inject_mode: InjectMode,
}

impl MessageAction {
    pub fn uuid(&self) -> &String {
        &self.uuid
    }
    pub fn name(&self) -> Option<&String> {
        self.name.as_ref()
    }
    pub fn message(&self) -> &Template {
        &self.message
    }
    pub fn values(&self) -> &Handlebars {
        &self.values
    }
    pub fn inject_mode(&self) -> &InjectMode {
        &self.inject_mode
    }

    fn render_value(&self, key: &String, template_context: &Context) -> Result<String, Error> {
        let mut writer = Vec::new();
        {
            try!(self.values.renderw(key, &template_context, &mut writer));
        }
        let string = try!(String::from_utf8(writer));
        Ok(string)
    }

    fn render_values(&self, template_context: &Context) -> Result<BTreeMap<String, String>, Error> {
        let mut rendered_values = BTreeMap::new();
        for (key, _) in self.values.get_templates() {
            let rendered_value = try!(self.render_value(key, &template_context));
            rendered_values.insert(key.to_string(), rendered_value);
        }
        Ok(rendered_values)
    }

    fn render_message(&self, state: &State, context: &BaseContext) -> Result<Message, Error> {
        let template_context = {
            use handlebars::Context;
            let context = RendererContext::new(state, context);
            Context::wraps(&context)
        };

        let mut rendered_values = try!(self.render_values(&template_context));
        MessageAction::extend_with_context_information(&mut rendered_values, state, context);
        let message = rendered_values.remove(MESSAGE)
                                     .expect(&format!("There is no '{}' key in the renderer \
                                                       key-value pairs",
                                                      MESSAGE));
        let name = self.name.as_ref().map(|name| name.borrow());
        let message = MessageBuilder::new(&self.uuid, message)
                          .name(name)
                          .values(rendered_values)
                          .build();
        Ok(message)
    }

    fn extend_with_context_information(values: &mut BTreeMap<String, String>,
                                       state: &State,
                                       context: &BaseContext) {
        values.insert(CONTEXT_UUID.to_string(),
                      context.uuid().to_hyphenated_string());
        values.insert(CONTEXT_LEN.to_string(), state.messages().len().to_string());
        if let Some(name) = context.name() {
            values.insert(CONTEXT_NAME.to_string(), name.to_string());
        }
    }

    fn execute(&self, _state: &State, _context: &BaseContext, responder: &mut ResponseSender) {
        match self.render_message(_state, _context) {
            Ok(message) => {
                let response = Alert {
                    message: message,
                    inject_mode: self.inject_mode.clone(),
                };
                responder.send_response(Response::Alert(response));
            }
            Err(error) => {
                error!("{}", error);
            }
        }
    }
}

impl From<MessageAction> for super::ActionType {
    fn from(action: MessageAction) -> super::ActionType {
        super::ActionType::Message(action)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum InjectMode {
    Log,
    Forward,
    Loopback,
}

impl Default for InjectMode {
    fn default() -> InjectMode {
        InjectMode::Log
    }
}

#[derive(Debug, Clone)]
pub struct Alert {
    message: Message,
    inject_mode: InjectMode,
}

impl Alert {
    pub fn message(&self) -> &Message {
        &self.message
    }
}

impl Action for MessageAction {
    fn on_opened(&self, _state: &State, _context: &BaseContext, responder: &mut ResponseSender) {
        if self.when.on_opened {
            trace!("MessageAction: on_opened()");
            self.execute(_state, _context, responder);
        }
    }

    fn on_closed(&self, _state: &State, _context: &BaseContext, responder: &mut ResponseSender) {
        if self.when.on_closed {
            trace!("MessageAction: on_closed()");
            self.execute(_state, _context, responder);
        }
    }
}
<commit_msg>MessageAction: make fields private<commit_after>use action::Action;
use context::base::BaseContext;
use dispatcher::Response;
use dispatcher::response::ResponseSender;
use message::{Message, MessageBuilder};

use handlebars::Context;
use std::borrow::Borrow;
use std::collections::BTreeMap;
use state::State;
use self::error::Error;
use self::renderer_context::RendererContext;

mod error;
mod renderer_context;
#[cfg(test)]
mod test;

pub const CONTEXT_UUID: &'static str = "context_uuid";
pub const CONTEXT_NAME: &'static str = "context_name";
pub const CONTEXT_LEN: &'static str = "context_len";
pub const MESSAGES: &'static str = "messages";
const MESSAGE: &'static str = "MESSAGE";

use handlebars::Template;
use handlebars::Handlebars;

use super::ActionType;
use super::ExecCondition;

mod deser;
mod builder;

pub use self::builder::MessageActionBuilder;

pub struct MessageAction {
    uuid: String,
    name: Option<String>,
    message: Template,
    values: Handlebars,
    when: ExecCondition,
    inject_mode: InjectMode,
}

impl MessageAction {
    pub fn uuid(&self) -> &String {
        &self.uuid
    }
    pub fn name(&self) -> Option<&String> {
        self.name.as_ref()
    }
    pub fn message(&self) -> &Template {
        &self.message
    }
    pub fn values(&self) -> &Handlebars {
        &self.values
    }
    pub fn inject_mode(&self) -> &InjectMode {
        &self.inject_mode
    }

    fn render_value(&self, key: &String, template_context: &Context) -> Result<String, Error> {
        let mut writer = Vec::new();
        {
            try!(self.values.renderw(key, &template_context, &mut writer));
        }
        let string = try!(String::from_utf8(writer));
        Ok(string)
    }

    fn render_values(&self, template_context: &Context) -> Result<BTreeMap<String, String>, Error> {
        let mut rendered_values = BTreeMap::new();
        for (key, _) in self.values.get_templates() {
            let rendered_value = try!(self.render_value(key, &template_context));
            rendered_values.insert(key.to_string(), rendered_value);
        }
        Ok(rendered_values)
    }

    fn render_message(&self, state: &State, context: &BaseContext) -> Result<Message, Error> {
        let template_context = {
            use handlebars::Context;
            let context = RendererContext::new(state, context);
            Context::wraps(&context)
        };

        let mut rendered_values = try!(self.render_values(&template_context));
        MessageAction::extend_with_context_information(&mut rendered_values, state, context);
        let message = rendered_values.remove(MESSAGE)
                                     .expect(&format!("There is no '{}' key in the renderer \
                                                       key-value pairs",
                                                      MESSAGE));
        let name = self.name.as_ref().map(|name| name.borrow());
        let message = MessageBuilder::new(&self.uuid, message)
                          .name(name)
                          .values(rendered_values)
                          .build();
        Ok(message)
    }

    fn extend_with_context_information(values: &mut BTreeMap<String, String>,
                                       state: &State,
                                       context: &BaseContext) {
        values.insert(CONTEXT_UUID.to_string(),
                      context.uuid().to_hyphenated_string());
        values.insert(CONTEXT_LEN.to_string(), state.messages().len().to_string());
        if let Some(name) = context.name() {
            values.insert(CONTEXT_NAME.to_string(), name.to_string());
        }
    }

    fn execute(&self, _state: &State, _context: &BaseContext, responder: &mut ResponseSender) {
        match self.render_message(_state, _context) {
            Ok(message) => {
                let response = Alert {
                    message: message,
                    inject_mode: self.inject_mode.clone(),
                };
                responder.send_response(Response::Alert(response));
            }
            Err(error) => {
                error!("{}", error);
            }
        }
    }
}

impl From<MessageAction> for super::ActionType {
    fn from(action: MessageAction) -> super::ActionType {
        super::ActionType::Message(action)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum InjectMode {
    Log,
    Forward,
    Loopback,
}

impl Default for InjectMode {
    fn default() -> InjectMode {
        InjectMode::Log
    }
}

#[derive(Debug, Clone)]
pub struct Alert {
    message: Message,
    inject_mode: InjectMode,
}

impl Alert {
    pub fn message(&self) -> &Message {
        &self.message
    }
}

impl Action for MessageAction {
    fn on_opened(&self, _state: &State, _context: &BaseContext, responder: &mut ResponseSender) {
        if self.when.on_opened {
            trace!("MessageAction: on_opened()");
            self.execute(_state, _context, responder);
        }
    }

    fn on_closed(&self, _state: &State, _context: &BaseContext, responder: &mut ResponseSender) {
        if self.when.on_closed {
            trace!("MessageAction: on_closed()");
            self.execute(_state, _context, responder);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;

use kvm_bindings::{CpuId, KVM_CPUID_FLAG_SIGNIFCANT_INDEX};

use crate::bit_helper::BitHelper;
use crate::cpu_leaf::*;
use crate::transformer::common::use_host_cpuid_function;

// Largest extended function. It has to be larger then 0x8000001d (Extended Cache Topology).
const LARGEST_EXTENDED_FN: u32 = 0x8000_001f;
// This value allows at most 64 logical threads within a package.
// See also the documentation for leaf_0x80000008::ecx::THREAD_ID_SIZE_BITRANGE
const THREAD_ID_MAX_SIZE: u32 = 7;
// This value means there is 1 node per processor.
// See also the documentation for leaf_0x8000001e::ecx::NODES_PER_PROCESSOR_BITRANGE.
const NODES_PER_PROCESSOR: u32 = 0;

pub fn update_structured_extended_entry(
    entry: &mut kvm_cpuid_entry2,
    _vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x7::index0::*;

    // according to the EPYC PPR, only the leaf 0x7 with index 0 contains the
    // structured extended feature identifiers
    if entry.index == 0 {
        // KVM sets this bit no matter what but this feature is not supported by hardware
        entry.edx.write_bit(edx::ARCH_CAPABILITIES_BITINDEX, false);
    }

    Ok(())
}

pub fn update_largest_extended_fn_entry(
    entry: &mut kvm_cpuid_entry2,
    _vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x80000000::*;

    // KVM sets the largest extended function to 0x80000000. Change it to 0x8000001f
    // Since we also use the leaf 0x8000001d (Extended Cache Topology).
    entry
        .eax
        .write_bits_in_range(&eax::LARGEST_EXTENDED_FN_BITRANGE, LARGEST_EXTENDED_FN);

    Ok(())
}

pub fn update_extended_feature_info_entry(
    entry: &mut kvm_cpuid_entry2,
    _vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x80000001::*;

    // set the Topology Extension bit since we use the Extended Cache Topology leaf
    entry.ecx.write_bit(ecx::TOPOEXT_INDEX, true);

    Ok(())
}

pub fn update_amd_features_entry(
    entry: &mut kvm_cpuid_entry2,
    vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x80000008::*;

    // We don't support more then 128 threads right now.
    // It's safe to put them all on the same processor.
    entry
        .ecx
        .write_bits_in_range(&ecx::THREAD_ID_SIZE_BITRANGE, THREAD_ID_MAX_SIZE)
        .write_bits_in_range(&ecx::NUM_THREADS_BITRANGE, u32::from(vm_spec.cpu_count - 1));

    Ok(())
}

pub fn update_extended_cache_topology_entry(
    entry: &mut kvm_cpuid_entry2,
    vm_spec: &VmSpec,
) -> Result<(), Error> {
    entry.flags |= KVM_CPUID_FLAG_SIGNIFCANT_INDEX;

    common::update_cache_parameters_entry(entry, vm_spec)
}

pub fn update_extended_apic_id_entry(
    entry: &mut kvm_cpuid_entry2,
    vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x8000001e::*;

    // When hyper-threading is enabled each pair of 2 consecutive logical CPUs
    // will have the same core id since they represent 2 threads in the same core.
    // For Example:
    // logical CPU 0 -> core id: 0
    // logical CPU 1 -> core id: 0
    // logical CPU 2 -> core id: 1
    // logical CPU 3 -> core id: 1
    let core_id = u32::from(vm_spec.cpu_index / vm_spec.cpus_per_core());

    entry
        .eax
        // the Extended APIC ID is the id of the current logical CPU
        .write_bits_in_range(
            &eax::EXTENDED_APIC_ID_BITRANGE,
            u32::from(vm_spec.cpu_index),
        );

    entry
        .ebx
        .write_bits_in_range(&ebx::CORE_ID_BITRANGE, core_id)
        .write_bits_in_range(
            &ebx::THREADS_PER_CORE_BITRANGE,
            u32::from(vm_spec.cpus_per_core() - 1),
        );

    entry
        .ecx
        .write_bits_in_range(&ecx::NODES_PER_PROCESSOR_BITRANGE, NODES_PER_PROCESSOR)
        // Put all the cpus in the same node.
        .write_bits_in_range(&ecx::NODE_ID_BITRANGE, 0);

    Ok(())
}

pub struct AmdCpuidTransformer {}

impl CpuidTransformer for AmdCpuidTransformer {
    fn process_cpuid(&self, cpuid: &mut CpuId, vm_spec: &VmSpec) -> Result<(), Error> {
        use_host_cpuid_function(cpuid, leaf_0x8000001e::LEAF_NUM, false)?;
        use_host_cpuid_function(cpuid, leaf_0x8000001d::LEAF_NUM, true)?;
        self.process_entries(cpuid, vm_spec)
    }

    fn entry_transformer_fn(&self, entry: &mut kvm_cpuid_entry2) -> Option<EntryTransformerFn> {
        match entry.function {
            leaf_0x1::LEAF_NUM => Some(common::update_feature_info_entry),
            leaf_0x7::LEAF_NUM => Some(amd::update_structured_extended_entry),
            leaf_0x80000000::LEAF_NUM => Some(amd::update_largest_extended_fn_entry),
            leaf_0x80000001::LEAF_NUM => Some(amd::update_extended_feature_info_entry),
            leaf_0x80000008::LEAF_NUM => Some(amd::update_amd_features_entry),
            leaf_0x8000001d::LEAF_NUM => Some(amd::update_extended_cache_topology_entry),
            leaf_0x8000001e::LEAF_NUM => Some(amd::update_extended_apic_id_entry),
            0x8000_0002..=0x8000_0004 => Some(common::update_brand_string_entry),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_update_structured_extended_entry() {
        use crate::cpu_leaf::leaf_0x7::index0::*;

        // Check that if index == 0 the entry is processed
        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: leaf_0x7::LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: *(0 as u32).write_bit(edx::ARCH_CAPABILITIES_BITINDEX, true),
            padding: [0, 0, 0],
        };
        assert!(update_structured_extended_entry(&mut entry, &vm_spec).is_ok());
        assert_eq!(entry.edx.read_bit(edx::ARCH_CAPABILITIES_BITINDEX), false);

        // Check that if index != 0 the entry is not processed
        entry.index = 1;
        entry.edx.write_bit(edx::ARCH_CAPABILITIES_BITINDEX, true);
        assert!(update_structured_extended_entry(&mut entry, &vm_spec).is_ok());
        assert_eq!(entry.edx.read_bit(edx::ARCH_CAPABILITIES_BITINDEX), true);
    }

    #[test]
    fn test_update_largest_extended_fn_entry() {
        use crate::cpu_leaf::leaf_0x80000000::*;

        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_largest_extended_fn_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(
            entry
                .eax
                .read_bits_in_range(&eax::LARGEST_EXTENDED_FN_BITRANGE),
            LARGEST_EXTENDED_FN
        );
    }

    #[test]
    fn test_update_extended_feature_info_entry() {
        use crate::cpu_leaf::leaf_0x80000001::*;

        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_extended_feature_info_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(entry.ecx.read_bit(ecx::TOPOEXT_INDEX), true);
    }

    fn check_update_amd_features_entry(cpu_count: u8, ht_enabled: bool) {
        use crate::cpu_leaf::leaf_0x80000008::*;

        let vm_spec = VmSpec::new(0, cpu_count, ht_enabled).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_amd_features_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(
            entry.ecx.read_bits_in_range(&ecx::NUM_THREADS_BITRANGE),
            u32::from(cpu_count - 1)
        );
        assert_eq!(
            entry.ecx.read_bits_in_range(&ecx::THREAD_ID_SIZE_BITRANGE),
            THREAD_ID_MAX_SIZE
        );
    }

    fn check_update_extended_apic_id_entry(
        cpu_id: u8,
        cpu_count: u8,
        ht_enabled: bool,
        expected_core_id: u32,
        expected_threads_per_core: u32,
    ) {
        use crate::cpu_leaf::leaf_0x8000001e::*;

        let vm_spec = VmSpec::new(cpu_id, cpu_count, ht_enabled).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_extended_apic_id_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(
            entry
                .eax
                .read_bits_in_range(&eax::EXTENDED_APIC_ID_BITRANGE),
            u32::from(cpu_id)
        );

        assert_eq!(
            entry.ebx.read_bits_in_range(&ebx::CORE_ID_BITRANGE),
            expected_core_id
        );
        assert_eq!(
            entry
                .ebx
                .read_bits_in_range(&ebx::THREADS_PER_CORE_BITRANGE),
            expected_threads_per_core
        );

        assert_eq!(
            entry
                .ecx
                .read_bits_in_range(&ecx::NODES_PER_PROCESSOR_BITRANGE),
            NODES_PER_PROCESSOR
        );
        assert_eq!(entry.ecx.read_bits_in_range(&ecx::NODE_ID_BITRANGE), 0);
    }

    #[test]
    fn test_update_extended_cache_topology_entry() {
        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: leaf_0x8000001d::LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_extended_cache_topology_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(entry.flags & KVM_CPUID_FLAG_SIGNIFCANT_INDEX, 1);
    }

    #[test]
    fn test_1vcpu_ht_off() {
        check_update_amd_features_entry(1, false);

        check_update_extended_apic_id_entry(0, 1, false, 0, 0);
    }

    #[test]
    fn test_1vcpu_ht_on() {
        check_update_amd_features_entry(1, true);

        check_update_extended_apic_id_entry(0, 1, true, 0, 0);
    }

    #[test]
    fn test_2vcpu_ht_off() {
        check_update_amd_features_entry(2, false);

        check_update_extended_apic_id_entry(0, 2, false, 0, 0);
        check_update_extended_apic_id_entry(1, 2, false, 1, 0);
    }

    #[test]
    fn test_2vcpu_ht_on() {
        check_update_amd_features_entry(2, true);

        check_update_extended_apic_id_entry(0, 2, true, 0, 1);
        check_update_extended_apic_id_entry(1, 2, true, 0, 1);
    }
}
<commit_msg>remove 0xB leaf on AMD<commit_after>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;

use kvm_bindings::{CpuId, KVM_CPUID_FLAG_SIGNIFCANT_INDEX};

use crate::bit_helper::BitHelper;
use crate::cpu_leaf::*;
use crate::transformer::common::use_host_cpuid_function;

// Largest extended function. It has to be larger then 0x8000001d (Extended Cache Topology).
const LARGEST_EXTENDED_FN: u32 = 0x8000_001f;
// This value allows at most 64 logical threads within a package.
// See also the documentation for leaf_0x80000008::ecx::THREAD_ID_SIZE_BITRANGE
const THREAD_ID_MAX_SIZE: u32 = 7;
// This value means there is 1 node per processor.
// See also the documentation for leaf_0x8000001e::ecx::NODES_PER_PROCESSOR_BITRANGE.
const NODES_PER_PROCESSOR: u32 = 0;

pub fn update_structured_extended_entry(
    entry: &mut kvm_cpuid_entry2,
    _vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x7::index0::*;

    // according to the EPYC PPR, only the leaf 0x7 with index 0 contains the
    // structured extended feature identifiers
    if entry.index == 0 {
        // KVM sets this bit no matter what but this feature is not supported by hardware
        entry.edx.write_bit(edx::ARCH_CAPABILITIES_BITINDEX, false);
    }

    Ok(())
}

pub fn update_largest_extended_fn_entry(
    entry: &mut kvm_cpuid_entry2,
    _vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x80000000::*;

    // KVM sets the largest extended function to 0x80000000. Change it to 0x8000001f
    // Since we also use the leaf 0x8000001d (Extended Cache Topology).
    entry
        .eax
        .write_bits_in_range(&eax::LARGEST_EXTENDED_FN_BITRANGE, LARGEST_EXTENDED_FN);

    Ok(())
}

pub fn update_extended_feature_info_entry(
    entry: &mut kvm_cpuid_entry2,
    _vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x80000001::*;

    // set the Topology Extension bit since we use the Extended Cache Topology leaf
    entry.ecx.write_bit(ecx::TOPOEXT_INDEX, true);

    Ok(())
}

pub fn update_amd_features_entry(
    entry: &mut kvm_cpuid_entry2,
    vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x80000008::*;

    // We don't support more then 128 threads right now.
    // It's safe to put them all on the same processor.
    entry
        .ecx
        .write_bits_in_range(&ecx::THREAD_ID_SIZE_BITRANGE, THREAD_ID_MAX_SIZE)
        .write_bits_in_range(&ecx::NUM_THREADS_BITRANGE, u32::from(vm_spec.cpu_count - 1));

    Ok(())
}

pub fn update_extended_cache_topology_entry(
    entry: &mut kvm_cpuid_entry2,
    vm_spec: &VmSpec,
) -> Result<(), Error> {
    entry.flags |= KVM_CPUID_FLAG_SIGNIFCANT_INDEX;

    common::update_cache_parameters_entry(entry, vm_spec)
}

pub fn update_extended_apic_id_entry(
    entry: &mut kvm_cpuid_entry2,
    vm_spec: &VmSpec,
) -> Result<(), Error> {
    use crate::cpu_leaf::leaf_0x8000001e::*;

    // When hyper-threading is enabled each pair of 2 consecutive logical CPUs
    // will have the same core id since they represent 2 threads in the same core.
    // For Example:
    // logical CPU 0 -> core id: 0
    // logical CPU 1 -> core id: 0
    // logical CPU 2 -> core id: 1
    // logical CPU 3 -> core id: 1
    let core_id = u32::from(vm_spec.cpu_index / vm_spec.cpus_per_core());

    entry
        .eax
        // the Extended APIC ID is the id of the current logical CPU
        .write_bits_in_range(
            &eax::EXTENDED_APIC_ID_BITRANGE,
            u32::from(vm_spec.cpu_index),
        );

    entry
        .ebx
        .write_bits_in_range(&ebx::CORE_ID_BITRANGE, core_id)
        .write_bits_in_range(
            &ebx::THREADS_PER_CORE_BITRANGE,
            u32::from(vm_spec.cpus_per_core() - 1),
        );

    entry
        .ecx
        .write_bits_in_range(&ecx::NODES_PER_PROCESSOR_BITRANGE, NODES_PER_PROCESSOR)
        // Put all the cpus in the same node.
        .write_bits_in_range(&ecx::NODE_ID_BITRANGE, 0);

    Ok(())
}

pub struct AmdCpuidTransformer {}

impl CpuidTransformer for AmdCpuidTransformer {
    fn process_cpuid(&self, cpuid: &mut CpuId, vm_spec: &VmSpec) -> Result<(), Error> {
        // Some versions of kernel may return the 0xB leaf for AMD even if this is an
        // Intel-specific leaf. Remove it.
        cpuid.retain(|entry| entry.function != leaf_0xb::LEAF_NUM);
        use_host_cpuid_function(cpuid, leaf_0x8000001e::LEAF_NUM, false)?;
        use_host_cpuid_function(cpuid, leaf_0x8000001d::LEAF_NUM, true)?;
        self.process_entries(cpuid, vm_spec)
    }

    fn entry_transformer_fn(&self, entry: &mut kvm_cpuid_entry2) -> Option<EntryTransformerFn> {
        match entry.function {
            leaf_0x1::LEAF_NUM => Some(common::update_feature_info_entry),
            leaf_0x7::LEAF_NUM => Some(amd::update_structured_extended_entry),
            leaf_0x80000000::LEAF_NUM => Some(amd::update_largest_extended_fn_entry),
            leaf_0x80000001::LEAF_NUM => Some(amd::update_extended_feature_info_entry),
            leaf_0x80000008::LEAF_NUM => Some(amd::update_amd_features_entry),
            leaf_0x8000001d::LEAF_NUM => Some(amd::update_extended_cache_topology_entry),
            leaf_0x8000001e::LEAF_NUM => Some(amd::update_extended_apic_id_entry),
            0x8000_0002..=0x8000_0004 => Some(common::update_brand_string_entry),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_process_cpuid() {
        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut cpuid = CpuId::new(0);
        let entry = kvm_cpuid_entry2 {
            function: leaf_0xb::LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };
        cpuid.push(entry).unwrap();

        assert!(AmdCpuidTransformer {}
            .process_cpuid(&mut cpuid, &vm_spec)
            .is_ok());
        assert!(!cpuid.as_slice().contains(&entry));
    }

    #[test]
    fn test_update_structured_extended_entry() {
        use crate::cpu_leaf::leaf_0x7::index0::*;

        // Check that if index == 0 the entry is processed
        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: leaf_0x7::LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: *(0 as u32).write_bit(edx::ARCH_CAPABILITIES_BITINDEX, true),
            padding: [0, 0, 0],
        };
        assert!(update_structured_extended_entry(&mut entry, &vm_spec).is_ok());
        assert_eq!(entry.edx.read_bit(edx::ARCH_CAPABILITIES_BITINDEX), false);

        // Check that if index != 0 the entry is not processed
        entry.index = 1;
        entry.edx.write_bit(edx::ARCH_CAPABILITIES_BITINDEX, true);
        assert!(update_structured_extended_entry(&mut entry, &vm_spec).is_ok());
        assert_eq!(entry.edx.read_bit(edx::ARCH_CAPABILITIES_BITINDEX), true);
    }

    #[test]
    fn test_update_largest_extended_fn_entry() {
        use crate::cpu_leaf::leaf_0x80000000::*;

        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_largest_extended_fn_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(
            entry
                .eax
                .read_bits_in_range(&eax::LARGEST_EXTENDED_FN_BITRANGE),
            LARGEST_EXTENDED_FN
        );
    }

    #[test]
    fn test_update_extended_feature_info_entry() {
        use crate::cpu_leaf::leaf_0x80000001::*;

        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_extended_feature_info_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(entry.ecx.read_bit(ecx::TOPOEXT_INDEX), true);
    }

    fn check_update_amd_features_entry(cpu_count: u8, ht_enabled: bool) {
        use crate::cpu_leaf::leaf_0x80000008::*;

        let vm_spec = VmSpec::new(0, cpu_count, ht_enabled).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_amd_features_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(
            entry.ecx.read_bits_in_range(&ecx::NUM_THREADS_BITRANGE),
            u32::from(cpu_count - 1)
        );
        assert_eq!(
            entry.ecx.read_bits_in_range(&ecx::THREAD_ID_SIZE_BITRANGE),
            THREAD_ID_MAX_SIZE
        );
    }

    fn check_update_extended_apic_id_entry(
        cpu_id: u8,
        cpu_count: u8,
        ht_enabled: bool,
        expected_core_id: u32,
        expected_threads_per_core: u32,
    ) {
        use crate::cpu_leaf::leaf_0x8000001e::*;

        let vm_spec = VmSpec::new(cpu_id, cpu_count, ht_enabled).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_extended_apic_id_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(
            entry
                .eax
                .read_bits_in_range(&eax::EXTENDED_APIC_ID_BITRANGE),
            u32::from(cpu_id)
        );

        assert_eq!(
            entry.ebx.read_bits_in_range(&ebx::CORE_ID_BITRANGE),
            expected_core_id
        );
        assert_eq!(
            entry
                .ebx
                .read_bits_in_range(&ebx::THREADS_PER_CORE_BITRANGE),
            expected_threads_per_core
        );

        assert_eq!(
            entry
                .ecx
                .read_bits_in_range(&ecx::NODES_PER_PROCESSOR_BITRANGE),
            NODES_PER_PROCESSOR
        );
        assert_eq!(entry.ecx.read_bits_in_range(&ecx::NODE_ID_BITRANGE), 0);
    }

    #[test]
    fn test_update_extended_cache_topology_entry() {
        let vm_spec = VmSpec::new(0, 1, false).expect("Error creating vm_spec");
        let mut entry = &mut kvm_cpuid_entry2 {
            function: leaf_0x8000001d::LEAF_NUM,
            index: 0,
            flags: 0,
            eax: 0,
            ebx: 0,
            ecx: 0,
            edx: 0,
            padding: [0, 0, 0],
        };

        assert!(update_extended_cache_topology_entry(&mut entry, &vm_spec).is_ok());

        assert_eq!(entry.flags & KVM_CPUID_FLAG_SIGNIFCANT_INDEX, 1);
    }

    #[test]
    fn test_1vcpu_ht_off() {
        check_update_amd_features_entry(1, false);

        check_update_extended_apic_id_entry(0, 1, false, 0, 0);
    }

    #[test]
    fn test_1vcpu_ht_on() {
        check_update_amd_features_entry(1, true);

        check_update_extended_apic_id_entry(0, 1, true, 0, 0);
    }

    #[test]
    fn test_2vcpu_ht_off() {
        check_update_amd_features_entry(2, false);

        check_update_extended_apic_id_entry(0, 2, false, 0, 0);
        check_update_extended_apic_id_entry(1, 2, false, 1, 0);
    }

    #[test]
    fn test_2vcpu_ht_on() {
        check_update_amd_features_entry(2, true);

        check_update_extended_apic_id_entry(0, 2, true, 0, 1);
        check_update_extended_apic_id_entry(1, 2, true, 0, 1);
    }
}
<|endoftext|>


<commit_before>//! A LinkedList that lets you insert elements from the front or the end of the list.
//! TODO: Add the ability to insert in-between a Node.
//! Lots of thanks to reem and Sharp for help with this from the #rust channel

use std::mem;

/// A Node is the basic unit in a LinkedList which stores the payload of type T.
#[derive(Debug,PartialEq)]
struct Node<T> {
    /// The payload of type T that is wrapped aorund a Node.
    pub payload: T,
    /// An Option of a boxed Node<T> that would point to the next Node<T>.
    next: Option<Box<Node<T>>>,
    /// An Option of a boxed Node<T> that would point to the previous Node<T>.
    prev: Option<Box<Node<T>>>,
}

impl<T> Node<T> {

    pub fn new(payload: T) -> Node<T> {
        //! Creates a instance of a `Node` with `next` and `prev` point to a Option::Node.
        Node::<T> {
            payload: payload,
            next: None,
            prev: None,
        }
    }

    pub fn payload(&self) -> Option<&T> {
        //! Access to a payload of Node.
        Some(&self.payload)
    }

    pub fn insert_after(&mut self, mut new: Box<Node<T>>) -> Option<Box<Node<T>>> {
        //! Takes a new boxed Node<T> and swaps it with any existing Node after the current one.
        //! We replace any Option::None or Option::Some with a new Option<Box<Node<T>>>
        mem::swap(&mut new.next, &mut self.next);
        mem::replace(&mut self.next, Some(new))
    }

    pub fn insert_before(&mut self, mut new: Box<Node<T>>) -> Option<Box<Node<T>>> {
        //! Takes a new boxed Node<T> and swaps it with any existing Node after the current one.
        //! We replace any Option::None or Option::Some with a new Option<Box<Node<T>>>
        mem::swap(&mut new.prev, &mut self.prev);
        mem::replace(&mut self.prev, Some(new))
    }

    pub fn remove_after(&mut self) -> Option<Box<Node<T>>> {
        //! Removes the Node after the current node.
        match self.next.take() {
            Some(mut next)  => {
                mem::swap(&mut self.next, &mut next.next);
                Some(next)
            },
            None => None,
        }
    }

    pub fn remove_before(&mut self) -> Option<Box<Node<T>>> {
        //! Removes the Node before the current node.
        match self.prev.take() {
            Some(mut prev)  => {
                mem::swap(&mut self.prev, &mut prev.prev);
                Some(prev)
            },
            None => None,
        }
    }

    pub fn next(&mut self) -> Option<&mut T> {
        //! Goes to the next Node and returns an Option to it.
        let Node {
            ref mut payload,
            ref mut prev,
            ref mut next,
        } = *self;

        match *next {
            Some(ref mut next) => {
                mem::swap(&mut next.payload, payload);

                let mut next = &mut **next;     // Cannot completely understand this bit yet.
                mem::swap(&mut next.next, &mut next.prev);

                mem::swap(&mut next.prev, prev);
            },
            None => return None,
        }
        mem::swap(prev, next);
        Some(payload)
    }

    pub fn prev(&mut self) -> Option<&mut T> {
        let Node {
            ref mut payload,
            ref mut prev,
            ref mut next,
        } = *self;

        match *prev {
            Some(ref mut prev) => {
                mem::swap(&mut prev.payload, payload);

                let mut prev = &mut **prev;     // Still don't understand this bit.
                mem::swap(&mut prev.prev, &mut prev.next);

                mem::swap(&mut prev.next, next);
            },
            None => return None,
        }
        mem::swap(prev, next);
        Some(payload)
    }
}

/// A LinkedList is the struct that you would need to use to create a list of elements.
/// It contains a pointer to the first and last node in the list represented as `head and `tail`
/// respectively.
pub struct LinkedList<T> {
    /// First element in the list.
    head: Option<Box<Node<T>>>,
    /// Last element in the list.
    tail: Option<Box<Node<T>>>,
}

impl<T> LinkedList<T> {
    pub fn new() -> LinkedList<T> {
        //! Create a new LinkedList to contain a payloads with type T.
        LinkedList::<T> {
            head: None,
            tail: None,
        }
    }
}

#[test]
fn node_new() {
    let expected = Node {
        payload: "test",
        next: None,
        prev: None,
    };

    let node: Node<&str> = Node::new("test");

    assert_eq!(expected.payload,  node.payload);
    assert_eq!(expected.next,     node.next);
    assert_eq!(expected.prev,     node.prev);
}

#[test]
fn node_payload() {
    let node: Node<&str> = Node::new("test");
    assert_eq!(&"test", node.payload().unwrap());
}

#[test]
fn node_insert_after() {
    let mut node: Node<&str> = Node::new("one");
    node.insert_after(Box::new(Node::new("two")));
    assert_eq!("two", node.next.unwrap().payload);

}

#[test]
fn node_insert_before() {
    let mut node: Node<&str> = Node::new("two");
    node.insert_before(Box::new(Node::new("one")));
    assert_eq!("one", node.prev.unwrap().payload);
}

#[test]
fn node_remove_next() {
    let mut node: Node<&str> = Node::new("one");
    node.insert_after(Box::new(Node::new("two")));
    node.remove_after();
}

#[test]
fn node_remove_prev() {
    let mut node: Node<&str> = Node::new("one");
    node.insert_before(Box::new(Node::new("two")));
    node.remove_before();
}

#[cfg(test)]
mod tests {
    use super::*;

}
<commit_msg>Added two tests for next() and prev()<commit_after>//! A LinkedList that lets you insert elements from the front or the end of the list.
//! TODO: Add the ability to insert in-between a Node.
//! Lots of thanks to reem and Sharp for help with this from the #rust channel

use std::mem;

/// A Node is the basic unit in a LinkedList which stores the payload of type T.
#[derive(Debug,PartialEq)]
struct Node<T> {
    /// The payload of type T that is wrapped aorund a Node.
    pub payload: T,
    /// An Option of a boxed Node<T> that would point to the next Node<T>.
    next: Option<Box<Node<T>>>,
    /// An Option of a boxed Node<T> that would point to the previous Node<T>.
    prev: Option<Box<Node<T>>>,
}

impl<T> Node<T> {

    pub fn new(payload: T) -> Node<T> {
        //! Creates a instance of a `Node` with `next` and `prev` point to a Option::Node.
        Node::<T> {
            payload: payload,
            next: None,
            prev: None,
        }
    }

    pub fn payload(&self) -> Option<&T> {
        //! Access to a payload of Node.
        Some(&self.payload)
    }

    pub fn insert_after(&mut self, mut new: Box<Node<T>>) -> Option<Box<Node<T>>> {
        //! Takes a new boxed Node<T> and swaps it with any existing Node after the current one.
        //! We replace any Option::None or Option::Some with a new Option<Box<Node<T>>>
        mem::swap(&mut new.next, &mut self.next);
        mem::replace(&mut self.next, Some(new))
    }

    pub fn insert_before(&mut self, mut new: Box<Node<T>>) -> Option<Box<Node<T>>> {
        //! Takes a new boxed Node<T> and swaps it with any existing Node after the current one.
        //! We replace any Option::None or Option::Some with a new Option<Box<Node<T>>>
        mem::swap(&mut new.prev, &mut self.prev);
        mem::replace(&mut self.prev, Some(new))
    }

    pub fn remove_after(&mut self) -> Option<Box<Node<T>>> {
        //! Removes the Node after the current node.
        match self.next.take() {
            Some(mut next)  => {
                mem::swap(&mut self.next, &mut next.next);
                Some(next)
            },
            None => None,
        }
    }

    pub fn remove_before(&mut self) -> Option<Box<Node<T>>> {
        //! Removes the Node before the current node.
        match self.prev.take() {
            Some(mut prev)  => {
                mem::swap(&mut self.prev, &mut prev.prev);
                Some(prev)
            },
            None => None,
        }
    }

    pub fn next(&mut self) -> Option<&mut T> {
        //! Goes to the next Node and returns an Option to it.
        let Node {
            ref mut payload,
            ref mut prev,
            ref mut next,
        } = *self;

        match *next {
            Some(ref mut next) => {
                mem::swap(&mut next.payload, payload);

                let mut next = &mut **next;     // Cannot completely understand this bit yet.
                mem::swap(&mut next.next, &mut next.prev);

                mem::swap(&mut next.prev, prev);
            },
            None => return None,
        }
        mem::swap(prev, next);
        Some(payload)
    }

    pub fn prev(&mut self) -> Option<&mut T> {
        let Node {
            ref mut payload,
            ref mut prev,
            ref mut next,
        } = *self;

        match *prev {
            Some(ref mut prev) => {
                mem::swap(&mut prev.payload, payload);

                let mut prev = &mut **prev;     // Still don't understand this bit.
                mem::swap(&mut prev.prev, &mut prev.next);

                mem::swap(&mut prev.next, next);
            },
            None => return None,
        }
        mem::swap(prev, next);
        Some(payload)
    }
}

/// A LinkedList is the struct that you would need to use to create a list of elements.
/// It contains a pointer to the first and last node in the list represented as `head and `tail`
/// respectively.
pub struct LinkedList<T> {
    /// First element in the list.
    head: Option<Box<Node<T>>>,
    /// Last element in the list.
    tail: Option<Box<Node<T>>>,
}

impl<T> LinkedList<T> {
    pub fn new() -> LinkedList<T> {
        //! Create a new LinkedList to contain a payloads with type T.
        LinkedList::<T> {
            head: None,
            tail: None,
        }
    }
}

#[test]
fn node_new() {
    let expected = Node {
        payload: "test",
        next: None,
        prev: None,
    };

    let node: Node<&str> = Node::new("test");

    assert_eq!(expected.payload,  node.payload);
    assert_eq!(expected.next,     node.next);
    assert_eq!(expected.prev,     node.prev);
}

#[test]
fn node_payload() {
    let node: Node<&str> = Node::new("test");
    assert_eq!(&"test", node.payload().unwrap());
}

#[test]
fn node_insert_after() {
    let mut node: Node<&str> = Node::new("one");
    node.insert_after(Box::new(Node::new("two")));
    assert_eq!("two", node.next.unwrap().payload);
}

#[test]
fn node_insert_before() {
    let mut node: Node<&str> = Node::new("two");
    node.insert_before(Box::new(Node::new("one")));
    assert_eq!("one", node.prev.unwrap().payload);
}

#[test]
fn node_remove_next() {
    let mut node: Node<&str> = Node::new("one");
    node.insert_after(Box::new(Node::new("two")));
    node.remove_after();
}

#[test]
fn node_remove_prev() {
    let mut node: Node<&str> = Node::new("one");
    node.insert_before(Box::new(Node::new("two")));
    node.remove_before();
}

#[test]
fn node_next() {
    let mut node: Node<&str> = Node::new("one");
    node.insert_after(Box::new(Node::new("two")));
    assert_eq!("two", *node.next().unwrap());
}

#[test]
fn node_prev() {
    let mut node: Node<&str> = Node::new("two");
    node.insert_before(Box::new(Node::new("one")));
    assert_eq!("one", *node.prev().unwrap());
}

#[cfg(test)]
mod tests {
    use super::*;

}
<|endoftext|>


<commit_before>
fn single_right() {

}<commit_msg>removed unused redirection file<commit_after><|endoftext|>