<commit_before>//! This module implements functionality for interacting with GitHub.

use oauth2::AccessToken;
use reqwest::{self, header};

use serde::de::DeserializeOwned;

use std::str;

use crate::util::errors::{cargo_err, internal, not_found, AppError, AppResult};
use reqwest::blocking::Client;

#[derive(Debug)]
pub struct GitHubClient {
    base_url: String,
    client: Option<Client>,
}

impl GitHubClient {
    pub fn new(client: Option<Client>, base_url: String) -> Self {
        Self { base_url, client }
    }

    pub fn current_user(&self, auth: &AccessToken) -> AppResult<GithubUser> {
        self.request("/user", auth)
    }

    pub fn org_by_name(&self, org_name: &str, auth: &AccessToken) -> AppResult<GitHubOrganization> {
        let url = format!("/orgs/{org_name}");
        self.request(&url, auth)
    }

    pub fn team_by_name(
        &self,
        org_name: &str,
        team_name: &str,
        auth: &AccessToken,
    ) -> AppResult<GitHubTeam> {
        let url = format!("/orgs/{org_name}/teams/{team_name}");
        self.request(&url, auth)
    }

    pub fn team_membership(
        &self,
        org_id: i32,
        team_id: i32,
        username: &str,
        auth: &AccessToken,
    ) -> AppResult<GitHubTeamMembership> {
        let url = format!("/organizations/{org_id}/team/{team_id}/memberships/{username}");
        self.request(&url, auth)
    }

    /// Does all the nonsense for sending a GET to Github. Doesn't handle parsing
    /// because custom error-code handling may be desirable. Use
    /// `parse_github_response` to handle the "common" processing of responses.
    pub fn request<T>(&self, url: &str, auth: &AccessToken) -> AppResult<T>
    where
        T: DeserializeOwned,
    {
        let url = format!("{}{}", self.base_url, url);
        info!("GITHUB HTTP: {url}");

        self.client()
            .get(&url)
            .header(header::ACCEPT, "application/vnd.github.v3+json")
            .header(header::AUTHORIZATION, format!("token {}", auth.secret()))
            .header(header::USER_AGENT, "crates.io (https://crates.io)")
            .send()?
            .error_for_status()
            .map_err(|e| handle_error_response(&e))?
            .json()
            .map_err(Into::into)
    }

    /// Returns a client for making HTTP requests to upload crate files.
    ///
    /// The client will go through a proxy if the application was configured via
    /// `TestApp::with_proxy()`.
    ///
    /// # Panics
    ///
    /// Panics if the application was not initialized with a client.  This should only occur in
    /// tests that were not properly initialized.
    fn client(&self) -> &Client {
        self.client
            .as_ref()
            .expect("No HTTP client is configured.  In tests, use `TestApp::with_proxy()`.")
    }
}

fn handle_error_response(error: &reqwest::Error) -> Box<dyn AppError> {
    use reqwest::StatusCode as Status;

    match error.status() {
        Some(Status::UNAUTHORIZED) | Some(Status::FORBIDDEN) => cargo_err(
            "It looks like you don't have permission \
             to query a necessary property from GitHub \
             to complete this request. \
             You may need to re-authenticate on \
             crates.io to grant permission to read \
             GitHub org memberships.",
        ),
        Some(Status::NOT_FOUND) => not_found(),
        _ => internal(&format_args!(
            "didn't get a 200 result from github: {}",
            error
        )),
    }
}

#[derive(Debug, Deserialize)]
pub struct GithubUser {
    pub avatar_url: Option<String>,
    pub email: Option<String>,
    pub id: i32,
    pub login: String,
    pub name: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct GitHubOrganization {
    pub id: i32, // unique GH id (needed for membership queries)
    pub avatar_url: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct GitHubTeam {
    pub id: i32,              // unique GH id (needed for membership queries)
    pub name: Option<String>, // Pretty name
    pub organization: GitHubOrganization,
}

#[derive(Debug, Deserialize)]
pub struct GitHubTeamMembership {
    pub state: String,
}

pub fn team_url(login: &str) -> String {
    let mut login_pieces = login.split(':');
    login_pieces.next();
    format!(
        "https://github.com/{}",
        login_pieces.next().expect("org failed"),
    )
}
<commit_msg>remove outdated comment<commit_after>//! This module implements functionality for interacting with GitHub.

use oauth2::AccessToken;
use reqwest::{self, header};

use serde::de::DeserializeOwned;

use std::str;

use crate::util::errors::{cargo_err, internal, not_found, AppError, AppResult};
use reqwest::blocking::Client;

#[derive(Debug)]
pub struct GitHubClient {
    base_url: String,
    client: Option<Client>,
}

impl GitHubClient {
    pub fn new(client: Option<Client>, base_url: String) -> Self {
        Self { base_url, client }
    }

    pub fn current_user(&self, auth: &AccessToken) -> AppResult<GithubUser> {
        self.request("/user", auth)
    }

    pub fn org_by_name(&self, org_name: &str, auth: &AccessToken) -> AppResult<GitHubOrganization> {
        let url = format!("/orgs/{org_name}");
        self.request(&url, auth)
    }

    pub fn team_by_name(
        &self,
        org_name: &str,
        team_name: &str,
        auth: &AccessToken,
    ) -> AppResult<GitHubTeam> {
        let url = format!("/orgs/{org_name}/teams/{team_name}");
        self.request(&url, auth)
    }

    pub fn team_membership(
        &self,
        org_id: i32,
        team_id: i32,
        username: &str,
        auth: &AccessToken,
    ) -> AppResult<GitHubTeamMembership> {
        let url = format!("/organizations/{org_id}/team/{team_id}/memberships/{username}");
        self.request(&url, auth)
    }

    /// Does all the nonsense for sending a GET to Github.
    pub fn request<T>(&self, url: &str, auth: &AccessToken) -> AppResult<T>
    where
        T: DeserializeOwned,
    {
        let url = format!("{}{}", self.base_url, url);
        info!("GITHUB HTTP: {url}");

        self.client()
            .get(&url)
            .header(header::ACCEPT, "application/vnd.github.v3+json")
            .header(header::AUTHORIZATION, format!("token {}", auth.secret()))
            .header(header::USER_AGENT, "crates.io (https://crates.io)")
            .send()?
            .error_for_status()
            .map_err(|e| handle_error_response(&e))?
            .json()
            .map_err(Into::into)
    }

    /// Returns a client for making HTTP requests to upload crate files.
    ///
    /// The client will go through a proxy if the application was configured via
    /// `TestApp::with_proxy()`.
    ///
    /// # Panics
    ///
    /// Panics if the application was not initialized with a client.  This should only occur in
    /// tests that were not properly initialized.
    fn client(&self) -> &Client {
        self.client
            .as_ref()
            .expect("No HTTP client is configured.  In tests, use `TestApp::with_proxy()`.")
    }
}

fn handle_error_response(error: &reqwest::Error) -> Box<dyn AppError> {
    use reqwest::StatusCode as Status;

    match error.status() {
        Some(Status::UNAUTHORIZED) | Some(Status::FORBIDDEN) => cargo_err(
            "It looks like you don't have permission \
             to query a necessary property from GitHub \
             to complete this request. \
             You may need to re-authenticate on \
             crates.io to grant permission to read \
             GitHub org memberships.",
        ),
        Some(Status::NOT_FOUND) => not_found(),
        _ => internal(&format_args!(
            "didn't get a 200 result from github: {}",
            error
        )),
    }
}

#[derive(Debug, Deserialize)]
pub struct GithubUser {
    pub avatar_url: Option<String>,
    pub email: Option<String>,
    pub id: i32,
    pub login: String,
    pub name: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct GitHubOrganization {
    pub id: i32, // unique GH id (needed for membership queries)
    pub avatar_url: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct GitHubTeam {
    pub id: i32,              // unique GH id (needed for membership queries)
    pub name: Option<String>, // Pretty name
    pub organization: GitHubOrganization,
}

#[derive(Debug, Deserialize)]
pub struct GitHubTeamMembership {
    pub state: String,
}

pub fn team_url(login: &str) -> String {
    let mut login_pieces = login.split(':');
    login_pieces.next();
    format!(
        "https://github.com/{}",
        login_pieces.next().expect("org failed"),
    )
}
<|endoftext|>


<commit_before>extern crate std;
extern crate hyper;
extern crate serde_json;

use std::string::String;
use std::io::Read;
use self::serde_json::Value;

#[derive(Debug)]
pub struct GithubError;

impl From<hyper::error::Error> for GithubError {
    fn from(_: hyper::error::Error) -> Self {
        GithubError
    }
}

impl From<serde_json::error::Error> for GithubError {
    fn from(_: serde_json::error::Error) -> Self {
        GithubError
    }
}

impl From<std::io::Error> for GithubError {
    fn from(_: std::io::Error) -> Self {
        GithubError
    }
}

pub fn get_display_text(user: &str,
                        repo: &str,
                        issue: &str,
                        printurl: bool)
                        -> Result<String, GithubError> {

    let url = format!("https://api.github.com/repos/{}/{}/issues/{}",
                      user,
                      repo,
                      issue);
    let client = hyper::client::Client::new();
    let mut resp = try!(client.get(&url)
        .header(hyper::header::UserAgent("idsbot".to_owned()))
        .send());

    let mut body = String::new();
    try!(resp.read_to_string(&mut body));
    let data: Value = try!(serde_json::from_str(&body));
    let obj = try!(data.as_object().ok_or(GithubError));

    let title = try!(try!(obj.get("title").ok_or(GithubError))
        .as_str()
        .ok_or(GithubError));
    let state = try!(try!(obj.get("state").ok_or(GithubError))
        .as_str()
        .ok_or(GithubError));
    let url = try!(try!(obj.get("html_url").ok_or(GithubError))
        .as_str()
        .ok_or(GithubError));
    let issuetype = match obj.get("pull_request") {
        Some(..) => "pull request",
        None => "issue",
    };

    if printurl {
        Ok(format!("[{}] [{}] {}\n{}", issuetype, state, title, url))
    } else {
        Ok(format!("[{}] [{}] {}", issuetype, state, title))
    }
}
<commit_msg>Use ? instead of try!<commit_after>extern crate std;
extern crate hyper;
extern crate serde_json;

use std::string::String;
use std::io::Read;
use self::serde_json::Value;

#[derive(Debug)]
pub struct GithubError;

impl From<hyper::error::Error> for GithubError {
    fn from(_: hyper::error::Error) -> Self {
        GithubError
    }
}

impl From<serde_json::error::Error> for GithubError {
    fn from(_: serde_json::error::Error) -> Self {
        GithubError
    }
}

impl From<std::io::Error> for GithubError {
    fn from(_: std::io::Error) -> Self {
        GithubError
    }
}

pub fn get_display_text(user: &str,
                        repo: &str,
                        issue: &str,
                        printurl: bool)
                        -> Result<String, GithubError> {

    let url = format!("https://api.github.com/repos/{}/{}/issues/{}",
                      user,
                      repo,
                      issue);
    let client = hyper::client::Client::new();
    let mut resp = client.get(&url)
        .header(hyper::header::UserAgent("idsbot".to_owned()))
        .send()?;

    let mut body = String::new();
    resp.read_to_string(&mut body)?;
    let data: Value = serde_json::from_str(&body)?;
    let obj = data.as_object().ok_or(GithubError)?;

    let title = obj.get("title").ok_or(GithubError)?
        .as_str()
        .ok_or(GithubError)?;
    let state = obj.get("state").ok_or(GithubError)?
        .as_str()
        .ok_or(GithubError)?;
    let url = obj.get("html_url").ok_or(GithubError)?
        .as_str()
        .ok_or(GithubError)?;
    let issuetype = match obj.get("pull_request") {
        Some(..) => "pull request",
        None => "issue",
    };

    if printurl {
        Ok(format!("[{}] [{}] {}\n{}", issuetype, state, title, url))
    } else {
        Ok(format!("[{}] [{}] {}", issuetype, state, title))
    }
}
<|endoftext|>


<commit_before>
use std::io::Read;  // Trait providing read_to_string()

use std::env;

use hyper;
use rustc_serialize;

use projects;
use groups;

pub const API_VERSION: u16 = 3;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub version: String,
    pub revision: String,
}

#[derive(Debug)]
pub struct Pagination {
    pub page: u16,
    pub per_page: u16,
}

#[derive(Debug)]
pub struct GitLab {
    scheme: String,
    domain: String,
    port: u16,
    private_token: String,
    client: hyper::Client,
    pagination: Pagination,
}

impl GitLab {

    pub fn new(scheme: &str, domain: &str, port: u16, private_token: &str) -> GitLab {
        GitLab {
            scheme: scheme.to_string(),
            domain: domain.to_string(),
            port:   port,
            private_token: private_token.to_string(),
            client: match env::var("HTTP_PROXY") {
                Ok(proxy) => {
                    println!("proxy: {:?}", proxy);
                    let proxy: Vec<&str> = proxy.trim_left_matches("http://").split(':').collect();
                    let hostname = proxy[0].to_string();
                    let port = proxy[1];

                    hyper::Client::with_http_proxy(hostname, port.parse().unwrap())
                },
                Err(_) => hyper::Client::new(),
            },
            pagination: Pagination {page: 1, per_page: 20},
        }
    }

    pub fn new_http(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("http", &domain, 80, &private_token)
    }

    pub fn new_https(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("https", &domain, 443, &private_token)
    }

    pub fn build_url(&self, command: &str) -> String {
        format!("{}://{}:{}/api/v{}/{}?private_token={}&page={}&per_page={}",
                                self.scheme,
                                self.domain,
                                self.port,
                                API_VERSION,
                                command,
                                self.private_token,
                                self.pagination.page,
                                self.pagination.per_page)
    }

    pub fn attempt_connection(&self) -> Result<hyper::client::Response, hyper::Error> {
        // self.get("version")
        let url = self.build_url("version");
        // Close connections after each GET.
        let res = self.client.get(&url).header(hyper::header::Connection::close()).send();

        res
    }

    pub fn set_pagination(&mut self, pagination: Pagination) {
        self.pagination = pagination;
    }

    pub fn get<T>(&self, command: &str) -> Result<T, rustc_serialize::json::DecoderError>
            where T: rustc_serialize::Decodable {

        let url = self.build_url(command);
        let mut res: hyper::client::Response =
                        self.client
                        .get(&url)
                        .header(hyper::header::Connection::close())
                        .send()
                        .unwrap();

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();

        rustc_serialize::json::decode(body.as_str())
    }

    pub fn version(&self) -> Result<Version, rustc_serialize::json::DecoderError> {
        self.get("version")
    }

    pub fn groups(&self) -> Result<groups::Groups, rustc_serialize::json::DecoderError> {
        self.get("groups")
    }

    pub fn projects(&self) -> Result<projects::Projects, rustc_serialize::json::DecoderError> {
        self.get("projects")
    }
}

#[cfg(test)]
mod tests {
    use std::env;
    use gitlab::GitLab;
    use gitlab::Pagination;

    #[test]
    fn list_groups() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        gl.set_pagination(Pagination{page: 1, per_page: 100});
        println!("gl: {:?}", gl);

        let groups = gl.groups().unwrap();
        println!("groups: {:?}", groups);
    }

    #[test]
    fn list_projects() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        gl.set_pagination(Pagination{page: 1, per_page: 100});
        println!("gl: {:?}", gl);

        println!("projects: {:?}", gl.projects().unwrap());
        // assert_eq!(gl.attempt_connection().unwrap().status, hyper::Ok);

    }
}
<commit_msg>Attempt to list all projects one at a time in a loop to detect errors.<commit_after>
use std::io::Read;  // Trait providing read_to_string()

use std::env;

use hyper;
use rustc_serialize;

use projects;
use groups;

pub const API_VERSION: u16 = 3;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub version: String,
    pub revision: String,
}

#[derive(Debug)]
pub struct Pagination {
    pub page: u16,
    pub per_page: u16,
}

#[derive(Debug)]
pub struct GitLab {
    scheme: String,
    domain: String,
    port: u16,
    private_token: String,
    client: hyper::Client,
    pagination: Pagination,
}

impl GitLab {

    pub fn new(scheme: &str, domain: &str, port: u16, private_token: &str) -> GitLab {
        GitLab {
            scheme: scheme.to_string(),
            domain: domain.to_string(),
            port:   port,
            private_token: private_token.to_string(),
            client: match env::var("HTTP_PROXY") {
                Ok(proxy) => {
                    println!("proxy: {:?}", proxy);
                    let proxy: Vec<&str> = proxy.trim_left_matches("http://").split(':').collect();
                    let hostname = proxy[0].to_string();
                    let port = proxy[1];

                    hyper::Client::with_http_proxy(hostname, port.parse().unwrap())
                },
                Err(_) => hyper::Client::new(),
            },
            pagination: Pagination {page: 1, per_page: 20},
        }
    }

    pub fn new_http(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("http", &domain, 80, &private_token)
    }

    pub fn new_https(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("https", &domain, 443, &private_token)
    }

    pub fn build_url(&self, command: &str) -> String {
        format!("{}://{}:{}/api/v{}/{}?private_token={}&page={}&per_page={}",
                                self.scheme,
                                self.domain,
                                self.port,
                                API_VERSION,
                                command,
                                self.private_token,
                                self.pagination.page,
                                self.pagination.per_page)
    }

    pub fn attempt_connection(&self) -> Result<hyper::client::Response, hyper::Error> {
        // self.get("version")
        let url = self.build_url("version");
        // Close connections after each GET.
        let res = self.client.get(&url).header(hyper::header::Connection::close()).send();

        res
    }

    pub fn set_pagination(&mut self, pagination: Pagination) {
        self.pagination = pagination;
    }

    pub fn get<T>(&self, command: &str) -> Result<T, rustc_serialize::json::DecoderError>
            where T: rustc_serialize::Decodable {

        let url = self.build_url(command);
        let mut res: hyper::client::Response =
                        self.client
                        .get(&url)
                        .header(hyper::header::Connection::close())
                        .send()
                        .unwrap();

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();

        rustc_serialize::json::decode(body.as_str())
    }

    pub fn version(&self) -> Result<Version, rustc_serialize::json::DecoderError> {
        self.get("version")
    }

    pub fn groups(&self) -> Result<groups::Groups, rustc_serialize::json::DecoderError> {
        self.get("groups")
    }

    pub fn projects(&self) -> Result<projects::Projects, rustc_serialize::json::DecoderError> {
        self.get("projects")
    }
}

#[cfg(test)]
mod tests {
    use std::env;
    use gitlab::GitLab;
    use gitlab::Pagination;

    #[test]
    fn list_groups() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        gl.set_pagination(Pagination{page: 1, per_page: 100});
        println!("gl: {:?}", gl);

        let groups = gl.groups().unwrap();
        println!("groups: {:?}", groups);
    }

    #[test]
    fn list_projects() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        for i in 1..82 {
            gl.set_pagination(Pagination{page: i, per_page: 1});
            println!("projects: {:?}", gl.projects().unwrap());
        }

        // assert_eq!(gl.attempt_connection().unwrap().status, hyper::Ok);

    }
}
<|endoftext|>


<commit_before>use crate::core_parser::CoreParser;
use crate::errors::ParseError;
use crate::models::HeaderProp;
use crate::parsing_utils::{le_f32, le_i32, le_u64};

/// Intermediate parsing structure for the header
#[derive(Debug, PartialEq)]
pub struct Header {
    pub major_version: i32,
    pub minor_version: i32,
    pub net_version: Option<i32>,
    pub game_type: String,
    pub properties: Vec<(String, HeaderProp)>,
}

impl Header {
    pub fn num_frames(&self) -> Option<i32> {
        self.properties
            .iter()
            .find(|&(key, _)| key == "NumFrames")
            .and_then(|&(_, ref prop)| {
                if let HeaderProp::Int(v) = *prop {
                    Some(v)
                } else {
                    None
                }
            })
    }

    pub fn max_channels(&self) -> Option<i32> {
        self.properties
            .iter()
            .find(|&(key, _)| key == "MaxChannels")
            .and_then(|&(_, ref prop)| {
                if let HeaderProp::Int(v) = *prop {
                    Some(v)
                } else {
                    None
                }
            })
    }
}

pub fn parse_header(rlp: &mut CoreParser) -> Result<Header, ParseError> {
    let major_version = rlp
        .take(4, le_i32)
        .map_err(|e| ParseError::ParseError("major version", rlp.bytes_read(), Box::new(e)))?;

    let minor_version = rlp
        .take(4, le_i32)
        .map_err(|e| ParseError::ParseError("minor version", rlp.bytes_read(), Box::new(e)))?;

    let net_version =
        if major_version > 865 && minor_version > 17 {
            Some(rlp.take(4, le_i32).map_err(|e| {
                ParseError::ParseError("net version", rlp.bytes_read(), Box::new(e))
            })?)
        } else {
            None
        };

    let game_type = rlp
        .parse_text()
        .map_err(|e| ParseError::ParseError("game type", rlp.bytes_read(), Box::new(e)))?;

    let properties = parse_rdict(rlp)
        .map_err(|e| ParseError::ParseError("header properties", rlp.bytes_read(), Box::new(e)))?;

    Ok(Header {
        major_version,
        minor_version,
        net_version,
        game_type,
        properties,
    })
}

fn parse_rdict(rlp: &mut CoreParser) -> Result<Vec<(String, HeaderProp)>, ParseError> {
    // Other the actual network data, the header property associative array is the hardest to parse.
    // The format is to:
    // - Read string
    // - If string is "None", we're done
    // - else we're dealing with a property, and the string just read is the key. Now deserialize the
    //   value.
    // The return type of this function is a key value vector because since there is no format
    // specification, we can't rule out duplicate keys. Possibly consider a multi-map in the future.

    let mut res: Vec<_> = Vec::new();
    loop {
        let key = rlp.parse_str()?;
        if key == "None" || key == "\0\0\0None" {
            break;
        }

        let val = match rlp.parse_str()? {
            "ArrayProperty" => array_property(rlp),
            "BoolProperty" => bool_property(rlp),
            "ByteProperty" => byte_property(rlp),
            "FloatProperty" => float_property(rlp),
            "IntProperty" => int_property(rlp),
            "NameProperty" => name_property(rlp),
            "QWordProperty" => qword_property(rlp),
            "StrProperty" => str_property(rlp),
            x => Err(ParseError::UnexpectedProperty(String::from(x))),
        }?;

        res.push((String::from(key), val));
    }

    Ok(res)
}

// Header properties are encoded in a pretty simple format, with some oddities. The first 64bits
// is data that can be discarded, some people think that the 64bits is the length of the data
// while others think that the first 32bits is the header length in bytes with the subsequent
// 32bits unknown. Doesn't matter to us, we throw it out anyways. The rest of the bytes are
// decoded property type specific.

fn byte_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    // It's unknown (to me at least) why the byte property has two strings in it.
    rlp.take(8, |_d| ())?;
    match rlp.parse_str()? {
        "OnlinePlatform_Steam" | "OnlinePlatform_PS4" => Ok(()),
        _ => rlp.parse_str().map(|_| ()),
    }?;
    Ok(HeaderProp::Byte)
}

fn str_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(8, |_d| ())?;
    Ok(HeaderProp::Str(rlp.parse_text()?))
}

fn name_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(8, |_d| ())?;
    Ok(HeaderProp::Name(rlp.parse_text()?))
}

fn int_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(12, |d| HeaderProp::Int(le_i32(&d[8..])))
}

fn bool_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(9, |d| HeaderProp::Bool(d[8] == 1))
}

fn float_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(12, |d| HeaderProp::Float(le_f32(&d[8..])))
}

fn qword_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(16, |d| HeaderProp::QWord(le_u64(&d[8..])))
}

fn array_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    let size = rlp.take(12, |d| le_i32(&d[8..]))?;
    let arr = CoreParser::repeat(size as usize, || parse_rdict(rlp))?;
    Ok(HeaderProp::Array(arr))
}

#[cfg(test)]
mod tests {
    use crate::core_parser::CoreParser;

    use super::*;

    #[test]
    fn rdict_no_elements() {
        let data = [0x05, 0x00, 0x00, 0x00, b'N', b'o', b'n', b'e', 0x00];
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, Vec::new());
    }

    #[test]
    fn rdict_one_element() {
        // dd skip=$((0x1269)) count=$((0x12a8 - 0x1269)) if=rumble.replay of=rdict_one.replay bs=1
        let data = include_bytes!("../assets/replays/partial/rdict_one.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(
            res,
            vec![(String::from("PlayerName"), HeaderProp::Str(String::from("comagoosie")))]
        );
    }

    #[test]
    fn rdict_one_int_element() {
        // dd skip=$((0x250)) count=$((0x284 - 0x250)) if=rumble.replay of=rdict_int.replay bs=1
        let data = include_bytes!("../assets/replays/partial/rdict_int.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("PlayerTeam"), HeaderProp::Int(0))]);
    }

    #[test]
    fn rdict_one_bool_element() {
        // dd skip=$((0xa0f)) count=$((0xa3b - 0xa0f)) if=rumble.replay of=rdict_bool.replay bs=1
        let data = include_bytes!("../assets/replays/partial/rdict_bool.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("bBot"), HeaderProp::Bool(false))]);
    }

    fn append_none(input: &[u8]) -> Vec<u8> {
        let append = [0x05, 0x00, 0x00, 0x00, b'N', b'o', b'n', b'e', 0x00];
        let mut v = Vec::new();
        v.extend_from_slice(input);
        v.extend_from_slice(&append);
        v
    }

    #[test]
    fn rdict_one_name_element() {
        // dd skip=$((0x1237)) count=$((0x1269 - 0x1237)) if=rumble.replay of=rdict_name.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_name.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(
            res,
            vec![(String::from("MatchType"), HeaderProp::Name(String::from("Online")))]
        );
    }

    #[test]
    fn rdict_one_float_element() {
        // dd skip=$((0x10a2)) count=$((0x10ce - 0x10a2)) if=rumble.replay of=rdict_float.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_float.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("RecordFPS"), HeaderProp::Float(30.0))]);
    }

    #[test]
    fn rdict_one_qword_element() {
        // dd skip=$((0x576)) count=$((0x5a5 - 0x576)) if=rumble.replay of=rdict_qword.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_qword.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(
            res,
            vec![(String::from("OnlineID"), HeaderProp::QWord(76561198101748375))]
        );
    }

    #[test]
    fn rdict_one_array_element() {
        // dd skip=$((0xab)) count=$((0x3f7 + 36)) if=rumble.replay of=rdict_array.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_array.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        let expected = vec![
            vec![
                (String::from("frame"), HeaderProp::Int(441)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("Cakeboss"))),
                (String::from("PlayerTeam"), HeaderProp::Int(1)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(1738)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("Sasha Kaun"))),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(3504)),
                (
                    String::from("PlayerName"),
                    HeaderProp::Str(String::from("SilentWarrior")),
                ),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(5058)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("jeffreyj1"))),
                (String::from("PlayerTeam"), HeaderProp::Int(1)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(5751)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("GOOSE LORD"))),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(6083)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("GOOSE LORD"))),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(7021)),
                (
                    String::from("PlayerName"),
                    HeaderProp::Str(String::from("SilentWarrior")),
                ),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
        ];
        assert_eq!(res, vec![(String::from("Goals"), HeaderProp::Array(expected))]);
    }

    #[test]
    fn rdict_one_byte_element() {
        // dd skip=$((0xdf0)) count=$((0xe41 - 0xdf0)) if=rumble.replay of=rdict_byte.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_byte.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("Platform"), HeaderProp::Byte)]);
    }

    #[test]
    fn rdict_ps4_online_id() {
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_ps4_online_id.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res.len(), 14);
        assert_eq!(res[0], (String::from("TeamSize"), HeaderProp::Int(3)));
    }
}
<commit_msg>Condense internal find header property<commit_after>use crate::core_parser::CoreParser;
use crate::errors::ParseError;
use crate::models::HeaderProp;
use crate::parsing_utils::{le_f32, le_i32, le_u64};

/// Intermediate parsing structure for the header
#[derive(Debug, PartialEq)]
pub struct Header {
    pub major_version: i32,
    pub minor_version: i32,
    pub net_version: Option<i32>,
    pub game_type: String,
    pub properties: Vec<(String, HeaderProp)>,
}

impl Header {
    pub fn num_frames(&self) -> Option<i32> {
        self.properties
            .iter()
            .find(|&(key, _)| key == "NumFrames")
            .and_then(|&(_, ref prop)| prop.as_i32())
    }

    pub fn max_channels(&self) -> Option<i32> {
        self.properties
            .iter()
            .find(|&(key, _)| key == "MaxChannels")
            .and_then(|&(_, ref prop)| prop.as_i32())
    }
}

pub fn parse_header(rlp: &mut CoreParser) -> Result<Header, ParseError> {
    let major_version = rlp
        .take(4, le_i32)
        .map_err(|e| ParseError::ParseError("major version", rlp.bytes_read(), Box::new(e)))?;

    let minor_version = rlp
        .take(4, le_i32)
        .map_err(|e| ParseError::ParseError("minor version", rlp.bytes_read(), Box::new(e)))?;

    let net_version =
        if major_version > 865 && minor_version > 17 {
            Some(rlp.take(4, le_i32).map_err(|e| {
                ParseError::ParseError("net version", rlp.bytes_read(), Box::new(e))
            })?)
        } else {
            None
        };

    let game_type = rlp
        .parse_text()
        .map_err(|e| ParseError::ParseError("game type", rlp.bytes_read(), Box::new(e)))?;

    let properties = parse_rdict(rlp)
        .map_err(|e| ParseError::ParseError("header properties", rlp.bytes_read(), Box::new(e)))?;

    Ok(Header {
        major_version,
        minor_version,
        net_version,
        game_type,
        properties,
    })
}

fn parse_rdict(rlp: &mut CoreParser) -> Result<Vec<(String, HeaderProp)>, ParseError> {
    // Other the actual network data, the header property associative array is the hardest to parse.
    // The format is to:
    // - Read string
    // - If string is "None", we're done
    // - else we're dealing with a property, and the string just read is the key. Now deserialize the
    //   value.
    // The return type of this function is a key value vector because since there is no format
    // specification, we can't rule out duplicate keys. Possibly consider a multi-map in the future.

    let mut res: Vec<_> = Vec::new();
    loop {
        let key = rlp.parse_str()?;
        if key == "None" || key == "\0\0\0None" {
            break;
        }

        let val = match rlp.parse_str()? {
            "ArrayProperty" => array_property(rlp),
            "BoolProperty" => bool_property(rlp),
            "ByteProperty" => byte_property(rlp),
            "FloatProperty" => float_property(rlp),
            "IntProperty" => int_property(rlp),
            "NameProperty" => name_property(rlp),
            "QWordProperty" => qword_property(rlp),
            "StrProperty" => str_property(rlp),
            x => Err(ParseError::UnexpectedProperty(String::from(x))),
        }?;

        res.push((String::from(key), val));
    }

    Ok(res)
}

// Header properties are encoded in a pretty simple format, with some oddities. The first 64bits
// is data that can be discarded, some people think that the 64bits is the length of the data
// while others think that the first 32bits is the header length in bytes with the subsequent
// 32bits unknown. Doesn't matter to us, we throw it out anyways. The rest of the bytes are
// decoded property type specific.

fn byte_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    // It's unknown (to me at least) why the byte property has two strings in it.
    rlp.take(8, |_d| ())?;
    match rlp.parse_str()? {
        "OnlinePlatform_Steam" | "OnlinePlatform_PS4" => Ok(()),
        _ => rlp.parse_str().map(|_| ()),
    }?;
    Ok(HeaderProp::Byte)
}

fn str_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(8, |_d| ())?;
    Ok(HeaderProp::Str(rlp.parse_text()?))
}

fn name_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(8, |_d| ())?;
    Ok(HeaderProp::Name(rlp.parse_text()?))
}

fn int_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(12, |d| HeaderProp::Int(le_i32(&d[8..])))
}

fn bool_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(9, |d| HeaderProp::Bool(d[8] == 1))
}

fn float_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(12, |d| HeaderProp::Float(le_f32(&d[8..])))
}

fn qword_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    rlp.take(16, |d| HeaderProp::QWord(le_u64(&d[8..])))
}

fn array_property(rlp: &mut CoreParser) -> Result<HeaderProp, ParseError> {
    let size = rlp.take(12, |d| le_i32(&d[8..]))?;
    let arr = CoreParser::repeat(size as usize, || parse_rdict(rlp))?;
    Ok(HeaderProp::Array(arr))
}

#[cfg(test)]
mod tests {
    use crate::core_parser::CoreParser;

    use super::*;

    #[test]
    fn rdict_no_elements() {
        let data = [0x05, 0x00, 0x00, 0x00, b'N', b'o', b'n', b'e', 0x00];
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, Vec::new());
    }

    #[test]
    fn rdict_one_element() {
        // dd skip=$((0x1269)) count=$((0x12a8 - 0x1269)) if=rumble.replay of=rdict_one.replay bs=1
        let data = include_bytes!("../assets/replays/partial/rdict_one.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(
            res,
            vec![(String::from("PlayerName"), HeaderProp::Str(String::from("comagoosie")))]
        );
    }

    #[test]
    fn rdict_one_int_element() {
        // dd skip=$((0x250)) count=$((0x284 - 0x250)) if=rumble.replay of=rdict_int.replay bs=1
        let data = include_bytes!("../assets/replays/partial/rdict_int.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("PlayerTeam"), HeaderProp::Int(0))]);
    }

    #[test]
    fn rdict_one_bool_element() {
        // dd skip=$((0xa0f)) count=$((0xa3b - 0xa0f)) if=rumble.replay of=rdict_bool.replay bs=1
        let data = include_bytes!("../assets/replays/partial/rdict_bool.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("bBot"), HeaderProp::Bool(false))]);
    }

    fn append_none(input: &[u8]) -> Vec<u8> {
        let append = [0x05, 0x00, 0x00, 0x00, b'N', b'o', b'n', b'e', 0x00];
        let mut v = Vec::new();
        v.extend_from_slice(input);
        v.extend_from_slice(&append);
        v
    }

    #[test]
    fn rdict_one_name_element() {
        // dd skip=$((0x1237)) count=$((0x1269 - 0x1237)) if=rumble.replay of=rdict_name.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_name.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(
            res,
            vec![(String::from("MatchType"), HeaderProp::Name(String::from("Online")))]
        );
    }

    #[test]
    fn rdict_one_float_element() {
        // dd skip=$((0x10a2)) count=$((0x10ce - 0x10a2)) if=rumble.replay of=rdict_float.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_float.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("RecordFPS"), HeaderProp::Float(30.0))]);
    }

    #[test]
    fn rdict_one_qword_element() {
        // dd skip=$((0x576)) count=$((0x5a5 - 0x576)) if=rumble.replay of=rdict_qword.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_qword.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(
            res,
            vec![(String::from("OnlineID"), HeaderProp::QWord(76561198101748375))]
        );
    }

    #[test]
    fn rdict_one_array_element() {
        // dd skip=$((0xab)) count=$((0x3f7 + 36)) if=rumble.replay of=rdict_array.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_array.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        let expected = vec![
            vec![
                (String::from("frame"), HeaderProp::Int(441)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("Cakeboss"))),
                (String::from("PlayerTeam"), HeaderProp::Int(1)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(1738)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("Sasha Kaun"))),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(3504)),
                (
                    String::from("PlayerName"),
                    HeaderProp::Str(String::from("SilentWarrior")),
                ),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(5058)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("jeffreyj1"))),
                (String::from("PlayerTeam"), HeaderProp::Int(1)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(5751)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("GOOSE LORD"))),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(6083)),
                (String::from("PlayerName"), HeaderProp::Str(String::from("GOOSE LORD"))),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
            vec![
                (String::from("frame"), HeaderProp::Int(7021)),
                (
                    String::from("PlayerName"),
                    HeaderProp::Str(String::from("SilentWarrior")),
                ),
                (String::from("PlayerTeam"), HeaderProp::Int(0)),
            ],
        ];
        assert_eq!(res, vec![(String::from("Goals"), HeaderProp::Array(expected))]);
    }

    #[test]
    fn rdict_one_byte_element() {
        // dd skip=$((0xdf0)) count=$((0xe41 - 0xdf0)) if=rumble.replay of=rdict_byte.replay bs=1
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_byte.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res, vec![(String::from("Platform"), HeaderProp::Byte)]);
    }

    #[test]
    fn rdict_ps4_online_id() {
        let data = append_none(include_bytes!(
            "../assets/replays/partial/rdict_ps4_online_id.replay"
        ));
        let mut parser = CoreParser::new(&data[..]);
        let res = parse_rdict(&mut parser).unwrap();
        assert_eq!(res.len(), 14);
        assert_eq!(res[0], (String::from("TeamSize"), HeaderProp::Int(3)));
    }
}
<|endoftext|>


<commit_before>pub trait Indented {
    fn indent_end(&self, indent_size: u8) -> Option<i32>;
    fn pop_indentation(&mut self, indent_size: u8);
}

impl Indented for String {
    fn indent_end(&self, indent_size: u8) -> Option<i32> {
        self.rfind(&" ".repeat(indent_size as usize)).map(|e| e as i32 + indent_size as i32)
    }

    fn pop_indentation(&mut self, indent_size: u8) {
        if let Some(end) = self.indent_end(indent_size) {
            let end = end as usize;
            self.drain((end - indent_size as usize)..end);
        }
    }
}
<commit_msg>Don't treat mid-line whitespace as indentation<commit_after>pub trait Indented {
    fn indent_end(&self, indent_size: u8) -> Option<i32>;
    fn pop_indentation(&mut self, indent_size: u8);
}

impl Indented for String {
    fn indent_end(&self, indent_size: u8) -> Option<i32> {
        let mut leading_spaces = self.len() - self.trim_left().len();
        if leading_spaces % (indent_size as usize) > 0 {
            leading_spaces -= leading_spaces % (indent_size as usize);
        }
        if leading_spaces > 0 {
            Some(leading_spaces as i32)
        } else {
            None
        }
    }

    fn pop_indentation(&mut self, indent_size: u8) {
        if let Some(end) = self.indent_end(indent_size) {
            let end = end as usize;
            self.drain((end - indent_size as usize)..end);
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::mem;

use ast::{Ast, Expr, Stmt, Param};
use ast::{ExprAssignType, ExprIdentType, ExprBinType, ExprUnType};
use ast::{ExprLitStrType, ExprLitIntType, ExprLitBoolType};
use ast::{StmtVarType, StmtExprType, StmtIfType, StmtReturnType};
use ast::{StmtLoopType, StmtWhileType, StmtBlockType, Function};
use ast::Expr::*;
use ast::Stmt::*;
use ast::visit::*;

use ctxt::*;

use ir;
use ir::*;
use ir::Instr::*;
use ir::Opnd::*;

use sym::BuiltinType;

pub fn generate<'a, 'ast>(ctxt: &Context<'a, 'ast>, fct: &'ast Function) {
    Generator::new(ctxt, fct).generate();
}

struct Generator<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    vreg: u32,
    result: Opnd,
    block_id: BlockId,
    ast_fct: &'ast Function,
    ir: Fct,
    var_map: HashMap<VarInfoId, VarId>
}

impl<'a, 'ast> Generator<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) -> Generator<'a, 'ast> {
        Generator {
            ctxt: ctxt,
            vreg: 0,
            result: OpndInt(0),
            block_id: BlockId(0),
            ast_fct: fct,
            ir: Fct::new(),
            var_map: HashMap::new(),
        }
    }

    fn generate(&mut self) {
        for p in &self.ast_fct.params {
            self.visit_param(p);
        }

        self.block_id = self.ir.add_block();
        self.visit_stmt(&self.ast_fct.block);

        self.ensure_return();

        let ir = mem::replace(&mut self.ir, Fct::new());
        ir::dump::dump(self.ctxt, &ir);

        self.ctxt.fct_info_mut(self.ast_fct.id, |fct| fct.ir = Some(ir));
    }

    fn add_stmt_var(&mut self, stmt: &'ast StmtVarType) {
        let var_id = self.ctxt.var(stmt.id, |ctxt_var, ctxt_var_id| {
            let ir_var_id = self.ir.add_var(stmt.name, ctxt_var.data_type, 0);
            self.var_map.insert(ctxt_var_id, ir_var_id);

            ir_var_id
        });

        if let Some(ref expr) = stmt.expr {
            self.visit_expr(expr);

            let src = self.result;
            self.add_instr_assign_var(var_id, src);
        }
    }

    fn add_stmt_expr(&mut self, stmt: &'ast StmtExprType) {
        self.visit_expr(&stmt.expr);
    }

    fn add_stmt_if(&mut self, stmt: &'ast StmtIfType) {
        let before_id = self.block_id;
        let then_id = self.ir.add_block();
        let after_id = self.ir.add_block();
        let else_id = if stmt.else_block.is_some() {
            self.ir.add_block()
        } else {
            after_id
        };

        self.visit_expr(&stmt.cond);
        let result = self.result;
        self.add_instr(Instr::test(result, then_id, else_id));

        self.block_id = then_id;
        self.visit_stmt(&stmt.then_block);
        self.add_instr(Instr::goto(after_id));
        self.block_mut().add_predecessor(before_id);

        if let Some(ref else_block) = stmt.else_block {
            self.block_id = else_id;
            self.visit_stmt(else_block);
            self.add_instr(Instr::goto(after_id));
            self.block_mut().add_predecessor(before_id);
        }

        self.block_id = after_id;
    }

    fn add_stmt_loop(&mut self, stmt: &'ast StmtLoopType) {
        let before_id = self.block_id;
        let loop_id = self.ir.add_block();
        let after_id = self.ir.add_block();

        self.add_instr(InstrGoto(loop_id));

        self.block_id = loop_id;
        self.visit_stmt(&stmt.block);
        self.add_instr(Instr::goto(loop_id));
        self.block_mut().add_predecessor(before_id);

        self.block_id = after_id;
    }

    fn add_stmt_while(&mut self, stmt: &'ast StmtWhileType) {

    }

    fn add_stmt_block(&mut self, stmt: &'ast StmtBlockType) {
        for s in &stmt.stmts {
            self.visit_stmt(s);
        }
    }

    fn add_stmt_return(&mut self, stmt: &'ast StmtReturnType) {
        if let Some(ref expr) = stmt.expr {
            self.visit_expr(expr);

            let result = self.result;
            self.add_instr(Instr::ret_with(result));
        } else {
            self.add_instr(Instr::ret());
        }
    }

    fn add_expr_lit_int(&mut self, lit: &'ast ExprLitIntType) {
        self.result = OpndInt(lit.value);
    }

    fn add_expr_lit_bool(&mut self, lit: &'ast ExprLitBoolType) {
        self.result = OpndBool(lit.value);
    }

    fn add_expr_lit_str(&mut self, lit: &'ast ExprLitStrType) {
        let dest = self.next_vreg();
        let instr = Instr::str(dest, lit.value.clone());
        self.add_instr(instr);

        self.result = dest;
    }

    fn add_expr_un(&mut self, expr: &'ast ExprUnType) {
        self.visit_expr(&expr.opnd);
        let src = self.result;

        let dest = self.next_vreg();
        let instr = Instr::un(dest, expr.op, src);
        self.add_instr(instr);
        self.result = dest;
    }

    fn add_expr_bin(&mut self, expr: &'ast ExprBinType) {
        self.visit_expr(&expr.lhs);
        let lhs = self.result;

        self.visit_expr(&expr.rhs);
        let rhs = self.result;

        let dest = self.next_vreg();
        let instr = Instr::bin(dest, lhs, expr.op, rhs);
        self.add_instr(instr);
        self.result = dest;
    }

    fn add_expr_ident(&mut self, expr: &'ast ExprIdentType) {
        let var_id = self.ctxt.var(expr.id, |_, ctxt_var_id| {
            *self.var_map.get(&ctxt_var_id).unwrap()
        });

        let ssa_index = self.ir.vars[var_id.0].ssa_index;
        self.result = OpndVar(var_id, ssa_index);
    }

    fn add_expr_assign(&mut self, expr: &'ast ExprAssignType) {
        let var_id = self.ctxt.var(expr.lhs.id(), |_, ctxt_var_id| {
            *self.var_map.get(&ctxt_var_id).unwrap()
        });

        self.visit_expr(&expr.rhs);
        let src = self.result;

        self.add_instr_assign_var(var_id, src);
    }

    fn ensure_return(&mut self) {
        if let Some(&InstrRet(_)) = self.block().last_instr() {
            // already ends with ret: do nothing
        } else if self.ctxt.fct_info(self.ast_fct.id, |fct| fct.return_type) == BuiltinType::Unit {
            self.add_instr(Instr::ret());
        }
    }

    fn add_instr_assign_var(&mut self, dest: VarId, src: Opnd) {
        let ssa_index = {
            let var = &mut self.ir.vars[dest.0];
            var.ssa_index += 1;

            var.ssa_index
        };

        self.add_instr_assign(OpndVar(dest, ssa_index), src);
    }

    fn add_instr_assign(&mut self, dest: Opnd, src: Opnd) {
        self.add_instr(Instr::assign(dest, src));
    }

    fn add_instr(&mut self, instr: Instr) {
        self.block_mut().add_instr(instr);
    }

    fn block_mut(&mut self) -> &mut Block {
        self.ir.block_mut(self.block_id)
    }

    fn block(&self) -> &Block {
        self.ir.block(self.block_id)
    }

    fn next_vreg(&mut self) -> Opnd {
        let vreg = self.vreg;

        self.vreg += 1;

        OpndReg(vreg)
    }
}

impl<'a, 'ast> Visitor<'ast> for Generator<'a, 'ast> {
    fn visit_param(&mut self, p: &'ast Param) {
        self.ctxt.var(p.id, |ctxt_var, ctxt_var_id| {
            let ir_var_id = self.ir.add_var(p.name, ctxt_var.data_type, 0);
            self.var_map.insert(ctxt_var_id, ir_var_id);
        });
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        match *s {
            StmtExpr(ref stmt) => self.add_stmt_expr(stmt),
            StmtIf(ref stmt) => self.add_stmt_if(stmt),
            StmtLoop(ref stmt) => self.add_stmt_loop(stmt),
            StmtWhile(ref stmt) => self.add_stmt_while(stmt),
            StmtBlock(ref stmt) => self.add_stmt_block(stmt),
            StmtVar(ref stmt) => self.add_stmt_var(stmt),
            StmtReturn(ref stmt) => self.add_stmt_return(stmt),
            _ => panic!("unsupported statement")
        }
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        match *e {
            ExprLitInt(ref expr) => self.add_expr_lit_int(expr),
            ExprLitBool(ref expr) => self.add_expr_lit_bool(expr),
            ExprLitStr(ref expr) => self.add_expr_lit_str(expr),
            ExprIdent(ref expr) => self.add_expr_ident(expr),
            ExprUn(ref expr) => self.add_expr_un(expr),
            ExprBin(ref expr) => self.add_expr_bin(expr),
            ExprAssign(ref expr) => self.add_expr_assign(expr),
            _ => panic!("unsupported expression")
        }
    }
}

#[cfg(test)]
mod tests {
    use ast::BinOp;
    use ctxt::*;
    use ir;
    use ir::*;
    use ir::Instr::*;
    use ir::Opnd::*;
    use test::parse;

    fn check_fct<F, T>(code: &'static str, fname: &'static str, f: F) -> T where F: FnOnce(&Context, &FctInfo) -> T {
        parse(code, |ctxt| {
            let name = ctxt.interner.intern(fname);
            let fct_info_id = ctxt.sym.borrow().get_function(name).unwrap();

            let fct = ctxt.fct_info_for_id(fct_info_id, |fct_info| fct_info.ast.unwrap());
            ir::gen::generate(ctxt, fct);

            ctxt.fct_info_for_id(fct_info_id, |fct_info| f(ctxt, fct_info))
        })
    }

    #[test]
    fn assign() {
        check_fct("fn f() { var x = 1; x = x + 2; x = x + 3; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(6, instrs.len());
            assert_eq!(Instr::assign(OpndVar(VarId(0), 1), OpndInt(1)), instrs[0]);

            let lhs = OpndVar(VarId(0), 1);
            let rhs = OpndInt(2);
            assert_eq!(Instr::bin(OpndReg(0), lhs, BinOp::Add, rhs), instrs[1]);

            assert_eq!(Instr::assign(OpndVar(VarId(0), 2), OpndReg(0)), instrs[2]);

            let lhs = OpndVar(VarId(0), 2);
            let rhs = OpndInt(3);
            assert_eq!(Instr::bin(OpndReg(1), lhs, BinOp::Add, rhs), instrs[3]);

            assert_eq!(Instr::assign(OpndVar(VarId(0), 3), OpndReg(1)), instrs[4]);

            assert_eq!(Instr::ret(), instrs[5]);
        });
    }

    #[test]
    fn return_value() {
        check_fct("fn f() -> int { return 1; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(1, instrs.len());
            assert_eq!(Instr::ret_with(OpndInt(1)), instrs[0]);
        });
    }

    #[test]
    fn return_without_value() {
        check_fct("fn f() { return; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(1, instrs.len());
            assert_eq!(Instr::ret(), instrs[0]);
        });
    }

    #[test]
    fn implicit_return() {
        check_fct("fn f() { }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(1, instrs.len());
            assert_eq!(Instr::ret(), instrs[0]);
        });
    }

    #[test]
    fn param() {
        check_fct("fn f(a: int) -> int { return a; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            let param = OpndVar(VarId(0), 0);
            assert_eq!(Instr::ret_with(param), instrs[0]);
        });
    }
}
<commit_msg>IR: test expression without side effect<commit_after>use std::collections::HashMap;
use std::mem;

use ast::{Ast, Expr, Stmt, Param};
use ast::{ExprAssignType, ExprIdentType, ExprBinType, ExprUnType};
use ast::{ExprLitStrType, ExprLitIntType, ExprLitBoolType};
use ast::{StmtVarType, StmtExprType, StmtIfType, StmtReturnType};
use ast::{StmtLoopType, StmtWhileType, StmtBlockType, Function};
use ast::Expr::*;
use ast::Stmt::*;
use ast::visit::*;

use ctxt::*;

use ir;
use ir::*;
use ir::Instr::*;
use ir::Opnd::*;

use sym::BuiltinType;

pub fn generate<'a, 'ast>(ctxt: &Context<'a, 'ast>, fct: &'ast Function) {
    Generator::new(ctxt, fct).generate();
}

struct Generator<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    vreg: u32,
    result: Opnd,
    block_id: BlockId,
    ast_fct: &'ast Function,
    ir: Fct,
    var_map: HashMap<VarInfoId, VarId>
}

impl<'a, 'ast> Generator<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) -> Generator<'a, 'ast> {
        Generator {
            ctxt: ctxt,
            vreg: 0,
            result: OpndInt(0),
            block_id: BlockId(0),
            ast_fct: fct,
            ir: Fct::new(),
            var_map: HashMap::new(),
        }
    }

    fn generate(&mut self) {
        for p in &self.ast_fct.params {
            self.visit_param(p);
        }

        self.block_id = self.ir.add_block();
        self.visit_stmt(&self.ast_fct.block);

        self.ensure_return();

        let ir = mem::replace(&mut self.ir, Fct::new());
        ir::dump::dump(self.ctxt, &ir);

        self.ctxt.fct_info_mut(self.ast_fct.id, |fct| fct.ir = Some(ir));
    }

    fn add_stmt_var(&mut self, stmt: &'ast StmtVarType) {
        let var_id = self.ctxt.var(stmt.id, |ctxt_var, ctxt_var_id| {
            let ir_var_id = self.ir.add_var(stmt.name, ctxt_var.data_type, 0);
            self.var_map.insert(ctxt_var_id, ir_var_id);

            ir_var_id
        });

        if let Some(ref expr) = stmt.expr {
            self.visit_expr(expr);

            let src = self.result;
            self.add_instr_assign_var(var_id, src);
        }
    }

    fn add_stmt_expr(&mut self, stmt: &'ast StmtExprType) {
        self.visit_expr(&stmt.expr);
    }

    fn add_stmt_if(&mut self, stmt: &'ast StmtIfType) {
        let before_id = self.block_id;
        let then_id = self.ir.add_block();
        let after_id = self.ir.add_block();
        let else_id = if stmt.else_block.is_some() {
            self.ir.add_block()
        } else {
            after_id
        };

        self.visit_expr(&stmt.cond);
        let result = self.result;
        self.add_instr(Instr::test(result, then_id, else_id));

        self.block_id = then_id;
        self.visit_stmt(&stmt.then_block);
        self.add_instr(Instr::goto(after_id));
        self.block_mut().add_predecessor(before_id);

        if let Some(ref else_block) = stmt.else_block {
            self.block_id = else_id;
            self.visit_stmt(else_block);
            self.add_instr(Instr::goto(after_id));
            self.block_mut().add_predecessor(before_id);
        }

        self.block_id = after_id;
    }

    fn add_stmt_loop(&mut self, stmt: &'ast StmtLoopType) {
        let before_id = self.block_id;
        let loop_id = self.ir.add_block();
        let after_id = self.ir.add_block();

        self.add_instr(InstrGoto(loop_id));

        self.block_id = loop_id;
        self.visit_stmt(&stmt.block);
        self.add_instr(Instr::goto(loop_id));
        self.block_mut().add_predecessor(before_id);

        self.block_id = after_id;
    }

    fn add_stmt_while(&mut self, stmt: &'ast StmtWhileType) {

    }

    fn add_stmt_block(&mut self, stmt: &'ast StmtBlockType) {
        for s in &stmt.stmts {
            self.visit_stmt(s);
        }
    }

    fn add_stmt_return(&mut self, stmt: &'ast StmtReturnType) {
        if let Some(ref expr) = stmt.expr {
            self.visit_expr(expr);

            let result = self.result;
            self.add_instr(Instr::ret_with(result));
        } else {
            self.add_instr(Instr::ret());
        }
    }

    fn add_expr_lit_int(&mut self, lit: &'ast ExprLitIntType) {
        self.result = OpndInt(lit.value);
    }

    fn add_expr_lit_bool(&mut self, lit: &'ast ExprLitBoolType) {
        self.result = OpndBool(lit.value);
    }

    fn add_expr_lit_str(&mut self, lit: &'ast ExprLitStrType) {
        let dest = self.next_vreg();
        let instr = Instr::str(dest, lit.value.clone());
        self.add_instr(instr);

        self.result = dest;
    }

    fn add_expr_un(&mut self, expr: &'ast ExprUnType) {
        self.visit_expr(&expr.opnd);
        let src = self.result;

        let dest = self.next_vreg();
        let instr = Instr::un(dest, expr.op, src);
        self.add_instr(instr);
        self.result = dest;
    }

    fn add_expr_bin(&mut self, expr: &'ast ExprBinType) {
        self.visit_expr(&expr.lhs);
        let lhs = self.result;

        self.visit_expr(&expr.rhs);
        let rhs = self.result;

        let dest = self.next_vreg();
        let instr = Instr::bin(dest, lhs, expr.op, rhs);
        self.add_instr(instr);
        self.result = dest;
    }

    fn add_expr_ident(&mut self, expr: &'ast ExprIdentType) {
        let var_id = self.ctxt.var(expr.id, |_, ctxt_var_id| {
            *self.var_map.get(&ctxt_var_id).unwrap()
        });

        let ssa_index = self.ir.vars[var_id.0].ssa_index;
        self.result = OpndVar(var_id, ssa_index);
    }

    fn add_expr_assign(&mut self, expr: &'ast ExprAssignType) {
        let var_id = self.ctxt.var(expr.lhs.id(), |_, ctxt_var_id| {
            *self.var_map.get(&ctxt_var_id).unwrap()
        });

        self.visit_expr(&expr.rhs);
        let src = self.result;

        self.add_instr_assign_var(var_id, src);
    }

    fn ensure_return(&mut self) {
        if let Some(&InstrRet(_)) = self.block().last_instr() {
            // already ends with ret: do nothing
        } else if self.ctxt.fct_info(self.ast_fct.id, |fct| fct.return_type) == BuiltinType::Unit {
            self.add_instr(Instr::ret());
        }
    }

    fn add_instr_assign_var(&mut self, dest: VarId, src: Opnd) {
        let ssa_index = {
            let var = &mut self.ir.vars[dest.0];
            var.ssa_index += 1;

            var.ssa_index
        };

        self.add_instr_assign(OpndVar(dest, ssa_index), src);
    }

    fn add_instr_assign(&mut self, dest: Opnd, src: Opnd) {
        self.add_instr(Instr::assign(dest, src));
    }

    fn add_instr(&mut self, instr: Instr) {
        self.block_mut().add_instr(instr);
    }

    fn block_mut(&mut self) -> &mut Block {
        self.ir.block_mut(self.block_id)
    }

    fn block(&self) -> &Block {
        self.ir.block(self.block_id)
    }

    fn next_vreg(&mut self) -> Opnd {
        let vreg = self.vreg;

        self.vreg += 1;

        OpndReg(vreg)
    }
}

impl<'a, 'ast> Visitor<'ast> for Generator<'a, 'ast> {
    fn visit_param(&mut self, p: &'ast Param) {
        self.ctxt.var(p.id, |ctxt_var, ctxt_var_id| {
            let ir_var_id = self.ir.add_var(p.name, ctxt_var.data_type, 0);
            self.var_map.insert(ctxt_var_id, ir_var_id);
        });
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        match *s {
            StmtExpr(ref stmt) => self.add_stmt_expr(stmt),
            StmtIf(ref stmt) => self.add_stmt_if(stmt),
            StmtLoop(ref stmt) => self.add_stmt_loop(stmt),
            StmtWhile(ref stmt) => self.add_stmt_while(stmt),
            StmtBlock(ref stmt) => self.add_stmt_block(stmt),
            StmtVar(ref stmt) => self.add_stmt_var(stmt),
            StmtReturn(ref stmt) => self.add_stmt_return(stmt),
            _ => panic!("unsupported statement")
        }
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        match *e {
            ExprLitInt(ref expr) => self.add_expr_lit_int(expr),
            ExprLitBool(ref expr) => self.add_expr_lit_bool(expr),
            ExprLitStr(ref expr) => self.add_expr_lit_str(expr),
            ExprIdent(ref expr) => self.add_expr_ident(expr),
            ExprUn(ref expr) => self.add_expr_un(expr),
            ExprBin(ref expr) => self.add_expr_bin(expr),
            ExprAssign(ref expr) => self.add_expr_assign(expr),
            _ => panic!("unsupported expression")
        }
    }
}

#[cfg(test)]
mod tests {
    use ast::BinOp;
    use ctxt::*;
    use ir;
    use ir::*;
    use ir::Instr::*;
    use ir::Opnd::*;
    use test::parse;

    fn check_fct<F, T>(code: &'static str, fname: &'static str, f: F) -> T where F: FnOnce(&Context, &FctInfo) -> T {
        parse(code, |ctxt| {
            let name = ctxt.interner.intern(fname);
            let fct_info_id = ctxt.sym.borrow().get_function(name).unwrap();

            let fct = ctxt.fct_info_for_id(fct_info_id, |fct_info| fct_info.ast.unwrap());
            ir::gen::generate(ctxt, fct);

            ctxt.fct_info_for_id(fct_info_id, |fct_info| f(ctxt, fct_info))
        })
    }

    #[test]
    fn assign() {
        check_fct("fn f() { var x = 1; x = x + 2; x = x + 3; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(6, instrs.len());
            assert_eq!(Instr::assign(OpndVar(VarId(0), 1), OpndInt(1)), instrs[0]);

            let lhs = OpndVar(VarId(0), 1);
            let rhs = OpndInt(2);
            assert_eq!(Instr::bin(OpndReg(0), lhs, BinOp::Add, rhs), instrs[1]);

            assert_eq!(Instr::assign(OpndVar(VarId(0), 2), OpndReg(0)), instrs[2]);

            let lhs = OpndVar(VarId(0), 2);
            let rhs = OpndInt(3);
            assert_eq!(Instr::bin(OpndReg(1), lhs, BinOp::Add, rhs), instrs[3]);

            assert_eq!(Instr::assign(OpndVar(VarId(0), 3), OpndReg(1)), instrs[4]);

            assert_eq!(Instr::ret(), instrs[5]);
        });
    }

    #[test]
    fn return_value() {
        check_fct("fn f() -> int { return 1; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(1, instrs.len());
            assert_eq!(Instr::ret_with(OpndInt(1)), instrs[0]);
        });
    }

    #[test]
    fn return_without_value() {
        check_fct("fn f() { return; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(1, instrs.len());
            assert_eq!(Instr::ret(), instrs[0]);
        });
    }

    #[test]
    fn implicit_return() {
        check_fct("fn f() { }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(1, instrs.len());
            assert_eq!(Instr::ret(), instrs[0]);
        });
    }

    #[test]
    fn param() {
        check_fct("fn f(a: int) -> int { return a; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            let param = OpndVar(VarId(0), 0);
            assert_eq!(Instr::ret_with(param), instrs[0]);
        });
    }

    #[test]
    fn expr() {
        check_fct("fn f() { 1; 2; 3; }", "f", |ctxt, fct| {
            let ir_fct = fct.ir.as_ref().unwrap();
            assert_eq!(1, ir_fct.blocks.len());

            let instrs = &ir_fct.blocks[0].instructions;
            assert_eq!(1, instrs.len());
            assert_eq!(Instr::ret(), instrs[0]);
        });
    }
}
<|endoftext|>


<commit_before>#![feature(alloc)]
#![feature(asm)]
#![feature(box_syntax)]
#![feature(core_simd)]
#![feature(core_slice_ext)]
#![feature(core_str_ext)]
#![feature(fnbox)]
#![feature(fundamental)]
#![feature(lang_items)]
#![feature(no_std)]
#![feature(rc_unique)]
#![feature(unboxed_closures)]
#![feature(unsafe_no_drop_flag)]
#![no_std]

extern crate alloc;

#[macro_use]
extern crate mopa;

use core::fmt;
use core::mem::size_of;
use core::mem::swap;
use core::ptr;

use common::pio::*;
use common::memory::*;

use drivers::keyboard::keyboard_init;
use drivers::mouse::mouse_init;
use drivers::pci::*;
use drivers::ps2::*;
use drivers::serial::*;

use graphics::bmp::*;

use programs::filemanager::*;
use programs::common::*;
use programs::session::*;

use schemes::file::*;
use schemes::http::*;
use schemes::memory::*;
use schemes::pci::*;
use schemes::random::*;

mod common {
    pub mod debug;
    pub mod elf;
    pub mod event;
    pub mod memory;
    pub mod pci;
    pub mod pio;
    pub mod random;
    pub mod resource;
    pub mod string;
    pub mod vec;
}

mod drivers {
    pub mod disk;
    pub mod keyboard;
    pub mod mouse;
    pub mod pci;
    pub mod ps2;
    pub mod serial;
}

mod filesystems {
    pub mod unfs;
}

mod graphics {
    pub mod bmp;
    pub mod color;
    pub mod display;
    pub mod point;
    pub mod size;
    pub mod window;
}

mod network {
    pub mod arp;
    pub mod common;
    pub mod ethernet;
    pub mod icmp;
    pub mod intel8254x;
    pub mod ipv4;
    pub mod rtl8139;
    pub mod tcp;
    pub mod udp;
}

mod programs {
    pub mod common;
    pub mod editor;
    pub mod executor;
    pub mod filemanager;
    pub mod session;
    pub mod viewer;
}

mod schemes {
    pub mod file;
    pub mod http;
    pub mod ide;
    pub mod memory;
    pub mod pci;
    pub mod random;
}

mod usb {
    pub mod xhci;
}

static mut session_ptr: *mut Session = 0 as *mut Session;
static mut events_ptr: *mut Vec<Event> = 0 as *mut Vec<Event>;

unsafe fn init(){
    serial_init();

    dd(size_of::<usize>() * 8);
    d(" bits");
    dl();

    page_init();
    cluster_init();

    session_ptr = alloc(size_of::<Session>()) as *mut Session;
    *session_ptr = Session::new();
    events_ptr = alloc(size_of::<Vec<Event>>()) as *mut Vec<Event>;
    *events_ptr = Vec::new();

    let session = &mut *session_ptr;

    session.items.insert(0, Rc::new(FileManager::new()));

    keyboard_init();
    mouse_init();

    session.modules.push(Rc::new(PS2));
    session.modules.push(Rc::new(Serial::new(0x3F8, 0x4)));

    pci_init(session);

    session.modules.push(Rc::new(FileScheme));
    session.modules.push(Rc::new(HTTPScheme));
    session.modules.push(Rc::new(MemoryScheme));
    session.modules.push(Rc::new(PCIScheme));
    session.modules.push(Rc::new(RandomScheme));

    URL::from_string("file:///background.bmp".to_string()).open_async(box |mut resource: Box<Resource>|{
        let mut vec: Vec<u8> = Vec::new();
        match resource.read_to_end(&mut vec) {
            Option::Some(0) => d("No background data\n"),
            Option::Some(len) => {
                (*session_ptr).display.background = BMP::from_data(vec.as_ptr() as usize);
            },
            Option::None => d("Background load error\n")
        }
    });
}

fn dr(reg: &str, value: u32){
    d(reg);
    d(": ");
    dh(value as usize);
    dl();
}

#[no_mangle]
//Take regs for kernel calls and exceptions
pub unsafe fn kernel(interrupt: u32, edi: u32, esi: u32, ebp: u32, esp: u32, ebx: u32, edx: u32, ecx: u32, eax: u32, eip: u32, eflags: u32) {
    let exception = |name: &str|{
        d(name);
        dl();

        dr("INT", interrupt);
        dr("EIP", eip);
        dr("EFLAGS", eflags);
        dr("EAX", eax);
        dr("EBX", ebx);
        dr("ECX", ecx);
        dr("EDX", edx);
        dr("EDI", edi);
        dr("ESI", esi);
        dr("EBP", ebp);
        dr("ESP", esp);

        asm!("cli");
        asm!("hlt");
        loop{}
    };

    match interrupt {
        0x20 => (), //timer
        0x21 => (*session_ptr).on_irq(0x1), //keyboard
        0x23 => (*session_ptr).on_irq(0x3), // serial 2 and 4
        0x24 => (*session_ptr).on_irq(0x4), // serial 1 and 3
        0x29 => (*session_ptr).on_irq(0x9), //pci
        0x2A => (*session_ptr).on_irq(0xA), //pci
        0x2B => (*session_ptr).on_irq(0xB), //pci
        0x2C => (*session_ptr).on_irq(0xC), //mouse
        0x2E => (*session_ptr).on_irq(0xE), //disk
        0x2F => (*session_ptr).on_irq(0xF), //disk
        0x80 => { // kernel calls
            match eax {
                0x1 => {
                    d("Open: ");
                    let url: &URL = &*(ebx as *const URL);
                    url.d();

                    let session = &mut *session_ptr;
                    ptr::write(ecx as *mut Box<Resource>, session.open(url));
                },
                0x2 => {
                    d("Open Async: ");
                    let url: &URL = &*(ebx as *const URL);
                    let callback: Box<FnBox(Box<Resource>)> = ptr::read(ecx as *const Box<FnBox(Box<Resource>)>);
                    unalloc(ecx as usize);
                    url.d();

                    let session = &mut *session_ptr;
                    session.open_async(url, callback);
                },
                0x3 => {
                    (*events_ptr).push(*(ebx as *const Event));
                }
                _ => {
                    d("System Call");
                    d(" EAX:");
                    dh(eax as usize);
                    d(" EBX:");
                    dh(ebx as usize);
                    d(" ECX:");
                    dh(ecx as usize);
                    d(" EDX:");
                    dh(edx as usize);
                    dl();
                }
            }
        },
        0xFF => { // main loop
            init();

            let session = &mut *session_ptr;
            (*events_ptr).push(RedrawEvent {
                redraw: REDRAW_ALL
            }.to_event());
            loop {
                asm!("cli");
                while (*events_ptr).len() > 0 {
                    session.on_poll();

                    let mut events_copy: Vec<Event> = Vec::new();
                    swap(&mut events_copy, &mut *(events_ptr));

                    asm!("sti");
                    //Can preempt event handling and redraw
                    session.handle_events(&mut events_copy);
                    session.redraw();
                    asm!("cli");
                }
                asm!("sti");
                asm!("hlt");
            }
        },
        0x0 => exception("Divide by zero exception"),
        0x1 => exception("Debug exception"),
        0x2 => exception("Non-maskable interrupt"),
        0x3 => exception("Breakpoint exception"),
        0x4 => exception("Overflow exception"),
        0x5 => exception("Bound range exceeded exception"),
        0x6 => exception("Invalid opcode exception"),
        0x7 => exception("Device not available exception"),
        0x8 => exception("Double fault"),
        0xA => exception("Invalid TSS exception"),
        0xB => exception("Segment not present exception"),
        0xC => exception("Stack-segment fault"),
        0xD => exception("General protection fault"),
        0xE => exception("Page fault"),
        0x10 => exception("x87 floating-point exception"),
        0x11 => exception("Alignment check exception"),
        0x12 => exception("Machine check exception"),
        0x13 => exception("SIMD floating-point exception"),
        0x14 => exception("Virtualization exception"),
        0x1E => exception("Security exception"),
        _ => {
            d("Interrupt: ");
            dh(interrupt as usize);
            dl();
        }
    }

    if interrupt >= 0x20 && interrupt < 0x30 {
        if interrupt >= 0x28 {
            outb(0xA0, 0x20);
        }

        outb(0x20, 0x20);
    }
}

#[lang = "panic_fmt"]
pub extern fn panic_fmt(fmt: fmt::Arguments, file: &'static str, line: u32) -> ! {
    d("PANIC: ");
    d(file);
    d(": ");
    dh(line as usize);
    dl();
    unsafe{
        asm!("cli");
        asm!("hlt");
    }
    loop{}
}

#[no_mangle]
pub extern "C" fn memcmp(a: *mut u8, b: *const u8, len: isize) -> isize {
    unsafe {
        let mut i = 0;
        while i < len {
            let c_a = *a.offset(i);
            let c_b = *b.offset(i);
            if c_a != c_b{
                return c_a as isize - c_b as isize;
            }
            i += 1;
        }
        return 0;
    }
}

#[no_mangle]
pub extern "C" fn memmove(dst: *mut u8, src: *const u8, len: isize){
    unsafe {
        if src < dst {
            let mut i = len;
            while i > 0 {
                i -= 1;
                *dst.offset(i) = *src.offset(i);
            }
        }else{
            let mut i = 0;
            while i < len {
                *dst.offset(i) = *src.offset(i);
                i += 1;
            }
        }
    }
}

#[no_mangle]
pub extern "C" fn memcpy(dst: *mut u8, src: *const u8, len: isize){
    unsafe {
        let mut i = 0;
        while i < len {
            *dst.offset(i) = *src.offset(i);
            i += 1;
        }
    }
}

#[no_mangle]
pub extern "C" fn memset(src: *mut u8, c: i32, len: isize) {
    unsafe {
        let mut i = 0;
        while i < len {
            *src.offset(i) = c as u8;
            i += 1;
        }
    }
}
<commit_msg>Preempting cleanup<commit_after>#![feature(alloc)]
#![feature(asm)]
#![feature(box_syntax)]
#![feature(core_simd)]
#![feature(core_slice_ext)]
#![feature(core_str_ext)]
#![feature(fnbox)]
#![feature(fundamental)]
#![feature(lang_items)]
#![feature(no_std)]
#![feature(rc_unique)]
#![feature(unboxed_closures)]
#![feature(unsafe_no_drop_flag)]
#![no_std]

extern crate alloc;

#[macro_use]
extern crate mopa;

use core::fmt;
use core::mem::size_of;
use core::mem::swap;
use core::ptr;

use common::pio::*;
use common::memory::*;

use drivers::keyboard::keyboard_init;
use drivers::mouse::mouse_init;
use drivers::pci::*;
use drivers::ps2::*;
use drivers::serial::*;

use graphics::bmp::*;

use programs::filemanager::*;
use programs::common::*;
use programs::session::*;

use schemes::file::*;
use schemes::http::*;
use schemes::memory::*;
use schemes::pci::*;
use schemes::random::*;

mod common {
    pub mod debug;
    pub mod elf;
    pub mod event;
    pub mod memory;
    pub mod pci;
    pub mod pio;
    pub mod random;
    pub mod resource;
    pub mod string;
    pub mod vec;
}

mod drivers {
    pub mod disk;
    pub mod keyboard;
    pub mod mouse;
    pub mod pci;
    pub mod ps2;
    pub mod serial;
}

mod filesystems {
    pub mod unfs;
}

mod graphics {
    pub mod bmp;
    pub mod color;
    pub mod display;
    pub mod point;
    pub mod size;
    pub mod window;
}

mod network {
    pub mod arp;
    pub mod common;
    pub mod ethernet;
    pub mod icmp;
    pub mod intel8254x;
    pub mod ipv4;
    pub mod rtl8139;
    pub mod tcp;
    pub mod udp;
}

mod programs {
    pub mod common;
    pub mod editor;
    pub mod executor;
    pub mod filemanager;
    pub mod session;
    pub mod viewer;
}

mod schemes {
    pub mod file;
    pub mod http;
    pub mod ide;
    pub mod memory;
    pub mod pci;
    pub mod random;
}

mod usb {
    pub mod xhci;
}

static mut session_ptr: *mut Session = 0 as *mut Session;
static mut events_ptr: *mut Vec<Event> = 0 as *mut Vec<Event>;

unsafe fn init(){
    serial_init();

    dd(size_of::<usize>() * 8);
    d(" bits");
    dl();

    page_init();
    cluster_init();

    session_ptr = alloc(size_of::<Session>()) as *mut Session;
    *session_ptr = Session::new();
    events_ptr = alloc(size_of::<Vec<Event>>()) as *mut Vec<Event>;
    *events_ptr = Vec::new();

    let session = &mut *session_ptr;

    session.items.insert(0, Rc::new(FileManager::new()));

    keyboard_init();
    mouse_init();

    session.modules.push(Rc::new(PS2));
    session.modules.push(Rc::new(Serial::new(0x3F8, 0x4)));

    pci_init(session);

    session.modules.push(Rc::new(FileScheme));
    session.modules.push(Rc::new(HTTPScheme));
    session.modules.push(Rc::new(MemoryScheme));
    session.modules.push(Rc::new(PCIScheme));
    session.modules.push(Rc::new(RandomScheme));

    URL::from_string("file:///background.bmp".to_string()).open_async(box |mut resource: Box<Resource>|{
        let mut vec: Vec<u8> = Vec::new();
        match resource.read_to_end(&mut vec) {
            Option::Some(0) => d("No background data\n"),
            Option::Some(len) => {
                (*session_ptr).display.background = BMP::from_data(vec.as_ptr() as usize);
            },
            Option::None => d("Background load error\n")
        }
    });
}

fn dr(reg: &str, value: u32){
    d(reg);
    d(": ");
    dh(value as usize);
    dl();
}

#[no_mangle]
//Take regs for kernel calls and exceptions
pub unsafe fn kernel(interrupt: u32, edi: u32, esi: u32, ebp: u32, esp: u32, ebx: u32, edx: u32, ecx: u32, eax: u32, eip: u32, eflags: u32) {
    let exception = |name: &str|{
        d(name);
        dl();

        dr("INT", interrupt);
        dr("EIP", eip);
        dr("EFLAGS", eflags);
        dr("EAX", eax);
        dr("EBX", ebx);
        dr("ECX", ecx);
        dr("EDX", edx);
        dr("EDI", edi);
        dr("ESI", esi);
        dr("EBP", ebp);
        dr("ESP", esp);

        asm!("cli");
        asm!("hlt");
        loop{}
    };

    match interrupt {
        0x20 => (), //timer
        0x21 => (*session_ptr).on_irq(0x1), //keyboard
        0x23 => (*session_ptr).on_irq(0x3), // serial 2 and 4
        0x24 => (*session_ptr).on_irq(0x4), // serial 1 and 3
        0x29 => (*session_ptr).on_irq(0x9), //pci
        0x2A => (*session_ptr).on_irq(0xA), //pci
        0x2B => (*session_ptr).on_irq(0xB), //pci
        0x2C => (*session_ptr).on_irq(0xC), //mouse
        0x2E => (*session_ptr).on_irq(0xE), //disk
        0x2F => (*session_ptr).on_irq(0xF), //disk
        0x80 => { // kernel calls
            match eax {
                0x1 => {
                    d("Open: ");
                    let url: &URL = &*(ebx as *const URL);
                    url.d();

                    let session = &mut *session_ptr;
                    ptr::write(ecx as *mut Box<Resource>, session.open(url));
                },
                0x2 => {
                    d("Open Async: ");
                    let url: &URL = &*(ebx as *const URL);
                    let callback: Box<FnBox(Box<Resource>)> = ptr::read(ecx as *const Box<FnBox(Box<Resource>)>);
                    unalloc(ecx as usize);
                    url.d();

                    let session = &mut *session_ptr;
                    session.open_async(url, callback);
                },
                0x3 => {
                    (*events_ptr).push(*(ebx as *const Event));
                }
                _ => {
                    d("System Call");
                    d(" EAX:");
                    dh(eax as usize);
                    d(" EBX:");
                    dh(ebx as usize);
                    d(" ECX:");
                    dh(ecx as usize);
                    d(" EDX:");
                    dh(edx as usize);
                    dl();
                }
            }
        },
        0xFF => { // main loop
            init();

            let session = &mut *session_ptr;
            (*events_ptr).push(RedrawEvent {
                redraw: REDRAW_ALL
            }.to_event());
            loop {
                //Polling must take place without interrupts to avoid interruption during I/O logic
                asm!("cli");
                    session.on_poll();
                loop {
                    //Copying the events must take place without interrupts to avoid pushes of events
                    asm!("cli");
                        let mut events_copy: Vec<Event> = Vec::new();
                        swap(&mut events_copy, &mut *events_ptr);

                    //Can preempt event handling and redraw
                    asm!("sti");
                        session.handle_events(&mut events_copy);
                        session.redraw();

                    //Checking for new events must take place without interrupts to avoid pushes of events
                    asm!("cli");
                        if (*events_ptr).len() == 0 {
                            break;
                        }
                }
                //On no new events, halt
                asm!("sti");
                asm!("hlt");
            }
        },
        0x0 => exception("Divide by zero exception"),
        0x1 => exception("Debug exception"),
        0x2 => exception("Non-maskable interrupt"),
        0x3 => exception("Breakpoint exception"),
        0x4 => exception("Overflow exception"),
        0x5 => exception("Bound range exceeded exception"),
        0x6 => exception("Invalid opcode exception"),
        0x7 => exception("Device not available exception"),
        0x8 => exception("Double fault"),
        0xA => exception("Invalid TSS exception"),
        0xB => exception("Segment not present exception"),
        0xC => exception("Stack-segment fault"),
        0xD => exception("General protection fault"),
        0xE => exception("Page fault"),
        0x10 => exception("x87 floating-point exception"),
        0x11 => exception("Alignment check exception"),
        0x12 => exception("Machine check exception"),
        0x13 => exception("SIMD floating-point exception"),
        0x14 => exception("Virtualization exception"),
        0x1E => exception("Security exception"),
        _ => {
            d("Interrupt: ");
            dh(interrupt as usize);
            dl();
        }
    }

    if interrupt >= 0x20 && interrupt < 0x30 {
        if interrupt >= 0x28 {
            outb(0xA0, 0x20);
        }

        outb(0x20, 0x20);
    }
}

#[lang = "panic_fmt"]
pub extern fn panic_fmt(fmt: fmt::Arguments, file: &'static str, line: u32) -> ! {
    d("PANIC: ");
    d(file);
    d(": ");
    dh(line as usize);
    dl();
    unsafe{
        asm!("cli");
        asm!("hlt");
    }
    loop{}
}

#[no_mangle]
pub extern "C" fn memcmp(a: *mut u8, b: *const u8, len: isize) -> isize {
    unsafe {
        let mut i = 0;
        while i < len {
            let c_a = *a.offset(i);
            let c_b = *b.offset(i);
            if c_a != c_b{
                return c_a as isize - c_b as isize;
            }
            i += 1;
        }
        return 0;
    }
}

#[no_mangle]
pub extern "C" fn memmove(dst: *mut u8, src: *const u8, len: isize){
    unsafe {
        if src < dst {
            let mut i = len;
            while i > 0 {
                i -= 1;
                *dst.offset(i) = *src.offset(i);
            }
        }else{
            let mut i = 0;
            while i < len {
                *dst.offset(i) = *src.offset(i);
                i += 1;
            }
        }
    }
}

#[no_mangle]
pub extern "C" fn memcpy(dst: *mut u8, src: *const u8, len: isize){
    unsafe {
        let mut i = 0;
        while i < len {
            *dst.offset(i) = *src.offset(i);
            i += 1;
        }
    }
}

#[no_mangle]
pub extern "C" fn memset(src: *mut u8, c: i32, len: isize) {
    unsafe {
        let mut i = 0;
        while i < len {
            *src.offset(i) = c as u8;
            i += 1;
        }
    }
}
<|endoftext|>