<commit_before>use std::sync::Arc;

use light_arena::Allocator;

use bsdf::{BxDF, FresnelBlend, TrowbridgeReitzDistribution, BSDF};
use interaction::SurfaceInteraction;
use material::{Material, TransportMode};
use paramset::TextureParams;
use spectrum::Spectrum;
use texture::Texture;

#[derive(Debug)]
pub struct SubstrateMaterial {
    kd: Arc<Texture<Spectrum> + Send + Sync>,
    ks: Arc<Texture<Spectrum> + Send + Sync>,
    nu: Arc<Texture<f32> + Send + Sync>,
    nv: Arc<Texture<f32> + Send + Sync>,
    remap_roughness: bool,
}

impl SubstrateMaterial {
    pub fn create(mp: &mut TextureParams) -> Arc<Material + Send + Sync> {
        let kd = mp.get_spectrum_texture("Kd", &Spectrum::grey(0.5));
        let ks = mp.get_spectrum_texture("Ks", &Spectrum::grey(0.5));
        let urough = mp.get_float_texture("uroughness", 0.1);
        let vrough = mp.get_float_texture("vroughness", 0.1);
        let remap_roughness = mp.find_bool("remaproughness", true);

        Arc::new(SubstrateMaterial {
                     kd,
                     ks,
                     nu: urough,
                     nv: vrough,
                     remap_roughness,
                 })
    }
}

impl Material for SubstrateMaterial {
    fn compute_scattering_functions<'a, 'b>(&self,
                                            si: &mut SurfaceInteraction<'a, 'b>,
                                            _mode: TransportMode,
                                            _allow_multiple_lobes: bool,
                                            arena: &'b Allocator) {
        let mut bxdfs = arena.alloc_slice::<&BxDF>(8);
        let mut i = 0;

        let d = self.kd.evaluate(si);
        let s = self.ks.evaluate(si);
        let mut roughu = self.nu.evaluate(si);
        let mut roughv = self.nv.evaluate(si);

        if !d.is_black() || !s.is_black() {
            if self.remap_roughness {
                roughu = TrowbridgeReitzDistribution::roughness_to_alpha(roughu);
                roughv = TrowbridgeReitzDistribution::roughness_to_alpha(roughv);
            }
            let distrib = arena <- TrowbridgeReitzDistribution::new(roughu, roughv);
            bxdfs[i] = arena <- FresnelBlend::new(d, s, distrib);
            i += 1;
        }

        unsafe {
            let ptr = bxdfs.as_mut_ptr();
            bxdfs = ::std::slice::from_raw_parts_mut(ptr, i);
        }

        let bsdf = BSDF::new(si, 1.0, bxdfs);
        si.bsdf = Some(Arc::new(bsdf));
    }
}
<commit_msg>Fix for Substrate material<commit_after>use std::sync::Arc;

use light_arena::Allocator;

use bsdf::{BxDF, FresnelBlend, TrowbridgeReitzDistribution, BSDF};
use interaction::SurfaceInteraction;
use material::{Material, TransportMode};
use paramset::TextureParams;
use spectrum::Spectrum;
use texture::Texture;

#[derive(Debug)]
pub struct SubstrateMaterial {
    kd: Arc<Texture<Spectrum> + Send + Sync>,
    ks: Arc<Texture<Spectrum> + Send + Sync>,
    nu: Arc<Texture<f32> + Send + Sync>,
    nv: Arc<Texture<f32> + Send + Sync>,
    remap_roughness: bool,
}

impl SubstrateMaterial {
    pub fn create(mp: &mut TextureParams) -> Arc<Material + Send + Sync> {
        let kd = mp.get_spectrum_texture("Kd", &Spectrum::grey(0.5));
        let ks = mp.get_spectrum_texture("Ks", &Spectrum::grey(0.5));
        let urough = mp.get_float_texture("uroughness", 0.1);
        let vrough = mp.get_float_texture("vroughness", 0.1);
        let remap_roughness = mp.find_bool("remaproughness", true);

        Arc::new(SubstrateMaterial {
                     kd,
                     ks,
                     nu: urough,
                     nv: vrough,
                     remap_roughness,
                 })
    }
}

impl Material for SubstrateMaterial {
    fn compute_scattering_functions<'a, 'b>(&self,
                                            si: &mut SurfaceInteraction<'a, 'b>,
                                            _mode: TransportMode,
                                            _allow_multiple_lobes: bool,
                                            arena: &'b Allocator) {
        let mut bxdfs = arena.alloc_slice::<&BxDF>(8);
        let mut i = 0;

        let d = self.kd.evaluate(si).clamp();
        let s = self.ks.evaluate(si).clamp();
        let mut roughu = self.nu.evaluate(si);
        let mut roughv = self.nv.evaluate(si);

        if !d.is_black() || !s.is_black() {
            if self.remap_roughness {
                roughu = TrowbridgeReitzDistribution::roughness_to_alpha(roughu);
                roughv = TrowbridgeReitzDistribution::roughness_to_alpha(roughv);
            }
            let distrib = arena <- TrowbridgeReitzDistribution::new(roughu, roughv);
            bxdfs[i] = arena <- FresnelBlend::new(s, d, distrib);
            i += 1;
        }

        unsafe {
            let ptr = bxdfs.as_mut_ptr();
            bxdfs = ::std::slice::from_raw_parts_mut(ptr, i);
        }

        let bsdf = BSDF::new(si, 1.0, bxdfs);
        si.bsdf = Some(Arc::new(bsdf));
    }
}
<|endoftext|>


<commit_before>use pact_matching::models::{Pact, Interaction, Request, OptionalBody, PactSpecification};
use pact_matching::models::matchingrules::*;
use pact_matching::models::generators::*;
use pact_matching::models::parse_query_string;

use std::collections::{BTreeMap, HashMap};
use log::{log, debug, warn};
use hyper::{Body, Response, Server, Error};
use hyper::header::ToStrError;
use hyper::service::service_fn_ok;
use futures::future::Future;
use futures::stream::Stream;

fn extract_path(uri: &hyper::Uri) -> String {
    uri.path_and_query()
        .map(|path_and_query| path_and_query.path())
        .unwrap_or("/")
        .into()
}

fn extract_query_string(uri: &hyper::Uri) -> Option<HashMap<String, Vec<String>>> {
    uri.path_and_query()
        .and_then(|path_and_query| path_and_query.query())
        .and_then(|query| parse_query_string(&query.into()))
}

fn extract_headers(headers: &hyper::HeaderMap) -> Result<Option<HashMap<String, String>>, ToStrError> {
    if headers.len() > 0 {
        let result: Result<HashMap<String, String>, ToStrError> = headers.keys()
            .map(|name| -> Result<(String, String), ToStrError> {
                let values = headers.get_all(name);
                let mut iter = values.iter();

                let first_value = iter.next().unwrap();

                if iter.next().is_some() {
                    warn!("Multiple headers associated with '{}', but only the first is matched", name);
                }

                Ok((name.as_str().into(), first_value.to_str()?.into()))
            })
            .collect();

        result.map(|map| Some(map))
    } else {
        Ok(None)
    }
}

pub fn extract_body(chunk: hyper::Chunk) -> OptionalBody {
    let bytes = chunk.into_bytes();
    if bytes.len() > 0 {
        OptionalBody::Present(bytes.to_vec())
    } else {
        OptionalBody::Empty
    }
}

fn hyper_request_to_pact_request(req: hyper::Request<Body>) -> impl Future<Item = Request, Error = hyper::Error> {
    let method = req.method().to_string();
    let path = extract_path(req.uri());
    let query = extract_query_string(req.uri());
    let headers = extract_headers(req.headers());

    req.into_body()
        .concat2()
        .and_then(|body_chunk| Ok(Request {
            method: req.method().to_string(),
            path: path,
            query: query,
            headers: headers,
            body: extract_body(body_chunk),
            matching_rules: MatchingRules::default(),
            generators: Generators::default()
        }))
}

pub fn start(
    id: String,
    pact: Pact,
    port: u16,
    shutdown: impl Future<Item = (), Error = ()>,
) -> (impl Future<Item = (), Error = Error>, u16) {
    let addr = ([0, 0, 0, 0], port).into();
    let server = Server::bind(&addr)
        .serve(|| {
            service_fn_ok(|req| {
                debug!("Creating pact request from hyper request");
                let req = hyper_request_to_pact_request(req);
                Response::new(Body::from("Hello World"))
            })
        });

    let port = server.local_addr().port();

    (server.with_graceful_shutdown(shutdown), port)
}
<commit_msg>Fix hyper_request_to_pact_request<commit_after>use pact_matching::models::{Pact, Interaction, Request, OptionalBody, PactSpecification};
use pact_matching::models::matchingrules::*;
use pact_matching::models::generators::*;
use pact_matching::models::parse_query_string;

use std::collections::{BTreeMap, HashMap};
use log::{log, debug, warn};
use hyper::{Body, Response, Server, Error};
use hyper::header::ToStrError;
use hyper::service::service_fn_ok;
use futures::future;
use futures::future::Future;
use futures::stream::Stream;

enum MockRequestError {
    InvalidHeaderEncoding,
    RequestBodyError
}

fn extract_path(uri: &hyper::Uri) -> String {
    uri.path_and_query()
        .map(|path_and_query| path_and_query.path())
        .unwrap_or("/")
        .into()
}

fn extract_query_string(uri: &hyper::Uri) -> Option<HashMap<String, Vec<String>>> {
    uri.path_and_query()
        .and_then(|path_and_query| path_and_query.query())
        .and_then(|query| parse_query_string(&query.into()))
}

fn extract_headers(headers: &hyper::HeaderMap) -> Result<Option<HashMap<String, String>>, MockRequestError> {
    if headers.len() > 0 {
        let result: Result<HashMap<String, String>, MockRequestError> = headers.keys()
            .map(|name| -> Result<(String, String), MockRequestError> {
                let values = headers.get_all(name);
                let mut iter = values.iter();

                let first_value = iter.next().unwrap();

                if iter.next().is_some() {
                    warn!("Multiple headers associated with '{}', but only the first is used", name);
                }

                Ok((
                    name.as_str().into(),
                    first_value.to_str()
                        .map_err(|err| MockRequestError::InvalidHeaderEncoding)?
                        .into()
                    )
                )
            })
            .collect();

        result.map(|map| Some(map))
    } else {
        Ok(None)
    }
}

pub fn extract_body(chunk: hyper::Chunk) -> OptionalBody {
    let bytes = chunk.into_bytes();
    if bytes.len() > 0 {
        OptionalBody::Present(bytes.to_vec())
    } else {
        OptionalBody::Empty
    }
}

fn hyper_request_to_pact_request(req: hyper::Request<Body>) -> impl Future<Item = Request, Error = MockRequestError> {
    let method = req.method().to_string();
    let path = extract_path(req.uri());
    let query = extract_query_string(req.uri());
    let headers = extract_headers(req.headers());

    future::done(headers)
        .and_then(move |headers| {
            req.into_body()
                .concat2()
                .map_err(|_| MockRequestError::RequestBodyError)
                .map(|body_chunk| (headers, body_chunk))
        })
        .and_then(|(headers, body_chunk)|
            Ok(Request {
                method: method,
                path: path,
                query: query,
                headers: headers,
                body: extract_body(body_chunk),
                matching_rules: MatchingRules::default(),
                generators: Generators::default()
            })
        )
}

pub fn start(
    id: String,
    pact: Pact,
    port: u16,
    shutdown: impl Future<Item = (), Error = ()>,
) -> (impl Future<Item = (), Error = Error>, u16) {
    let addr = ([0, 0, 0, 0], port).into();
    let server = Server::bind(&addr)
        .serve(|| {
            service_fn_ok(|req| {
                debug!("Creating pact request from hyper request");
                let req = hyper_request_to_pact_request(req);
                Response::new(Body::from("Hello World"))
            })
        });

    let port = server.local_addr().port();

    (server.with_graceful_shutdown(shutdown), port)
}
<|endoftext|>


<commit_before>use crate::sys::config;
use crate::sys::taskserver::Control;
use super::SystemModule;
use anyhow::{Result, anyhow, ensure};
use chrono::NaiveTime;
use log::info;
use serde::{Serialize, Deserialize};
use tokio::process::Command;

#[derive(Clone, Serialize, Deserialize)]
struct HealthConfig {
    enabled: bool,
    debug_exec_once: bool,
}

pub struct Health {
    config: HealthConfig,
    wakeup_list: Vec<NaiveTime>,
}

impl Health {
    pub fn new(wakeup_list: Vec<NaiveTime>) -> Result<Self> {
        info!("[health] initialize");

        let jsobj = config::get_object(&["health"])
            .map_or(Err(anyhow!("Config not found: health")), Ok)?;
        let config: HealthConfig = serde_json::from_value(jsobj)?;

        Ok(Health {
            config,
            wakeup_list,
        })
    }

    async fn health_task(&mut self, ctrl: &Control) -> Result<()> {
        let cpu_info = get_cpu_info().await;
        info!("{:?}", cpu_info);
        let mem_info = get_mem_info().await;
        info!("{:?}", mem_info?);
        let disk_info = get_disk_info().await;
        info!("{:?}", disk_info?);
        let cpu_temp = get_cpu_temp().await;
        info!("{:?}", cpu_temp?);

        Ok(())
    }

    async fn health_task_entry(ctrl: Control) -> Result<()> {
        let mut health = ctrl.sysmods().health.write().await;
        health.health_task(&ctrl).await
    }
}

impl SystemModule for Health {
    fn on_start(&self, ctrl: &Control) {
        info!("[health] on_start");
        if self.config.enabled {
            if self.config.debug_exec_once {
                ctrl.spawn_oneshot_task(
                    "health_check",
                    Health::health_task_entry);
            }
            else {
                ctrl.spawn_periodic_task(
                    "health_check",
                    &self.wakeup_list,
                    Health::health_task_entry);
            }
        }
    }
}

#[derive(Debug, Clone)]
struct CpuInfo {
    cpu_percent_total: f64,
    cpu_percent_list: Vec<f64>,
}

async fn get_cpu_info() -> Result<CpuInfo> {
    let buf = tokio::fs::read("/proc/stat").await?;
    let text = String::from_utf8_lossy(&buf);

    // See `man proc`
    // user   (1) Time spent in user mode.
    // nice   (2) Time spent in user mode with low priority (nice).
    // system (3) Time spent in system mode.
    // idle   (4) Time spent in the idle task.  This value should be USER_HZ times the second entry in the /proc/uptime pseudo-file.
    // iowait (since Linux 2.5.41)
    //        (5) Time waiting for I/O to complete.  This value is not reliable, for the following reasons:
    // irq (since Linux 2.6.0-test4)
    //        (6) Time servicing interrupts.
    // softirq (since Linux 2.6.0-test4)
    //        (7) Time servicing softirqs.
    // steal (since Linux 2.6.11)
    //        (8)  Stolen  time,  which  is the time spent in other operating systems when running in a virtualized environment
    // guest (since Linux 2.6.24)
    //        (9) Time spent running a virtual CPU for guest operating systems under the control of the Linux kernel.
    // guest_nice (since Linux 2.6.33)
    //        (10)  Time spent running a niced guest (virtual CPU for guest operating systems under the
    //        control of the Linux kernel).
    let mut cpu_percent_total = None;
    let mut cpu_percent_list = vec![];
    for line in text.lines() {
        let mut name = None;
        let mut user = None;
        let mut nice = None;
        let mut system = None;
        let mut idle = None;
        for (col_no, elem) in line.split_ascii_whitespace().enumerate() {
            match col_no {
                0 => name = Some(elem),
                1 => user = Some(elem),
                2 => nice = Some(elem),
                3 => system = Some(elem),
                4 => idle = Some(elem),
                _ => (),
            }
        }
        // cpu or cpu%d の行を取り出す
        if name.is_none() || !name.unwrap().starts_with("cpu") {
            continue;
        }

        let user: u64 = user.ok_or(anyhow!("parse error"))?.parse()?;
        let nice: u64 = nice.ok_or(anyhow!("parse error"))?.parse()?;
        let system: u64 = system.ok_or(anyhow!("parse error"))?.parse()?;
        let idle: u64 = idle.ok_or(anyhow!("parse error"))?.parse()?;
        let total = user + nice + system + idle;
        let value = (total - idle) as f64 / total as f64;
        if name == Some("cpu") {
            cpu_percent_total = Some(value);
        }
        else {
            cpu_percent_list.push(value);
        }
    }

    ensure!(cpu_percent_total.is_some());
    ensure!(!cpu_percent_list.is_empty());
    let cpu_percent_total = cpu_percent_total.unwrap();
    Ok(CpuInfo {
        cpu_percent_total,
        cpu_percent_list,
    })
}

#[derive(Debug, Clone, Copy)]
struct MemInfo {
    total_mib: f64,
    avail_mib: f64,
}

async fn get_mem_info() -> Result<MemInfo> {
    let mut cmd = Command::new("free");
	let output = cmd.output().await?;
    ensure!(output.status.success(), "free command failed");

    // sample
    //                total        used        free      shared  buff/cache   available
    // Mem:        13034888     4119272     5561960          68     3353656     8609008
    // Swap:        4194304           0     4194304
    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut total = None;
    let mut avail = None;
    for (line_no, line) in stdout.lines().enumerate() {
        if line_no != 1 {
            continue;
        }
        for (col_no, elem) in line.split_ascii_whitespace().enumerate() {
            match col_no {
                1 => total = Some(elem),
                6 => avail = Some(elem),
                _ => (),
            }
        }
        break;
    }
    let total = total.ok_or(anyhow!("parse error"))?;
    let avail = avail.ok_or(anyhow!("parse error"))?;
    let total_mib = total.parse::<u64>()? as f64 / 1024.0;
    let avail_mib = avail.parse::<u64>()? as f64 / 1024.0;

    Ok(MemInfo { total_mib, avail_mib })
}

#[derive(Debug, Clone, Copy)]
struct DiskInfo {
    total_gib: f64,
    avail_gib: f64,
}

async fn get_disk_info() -> Result<DiskInfo> {
    let mut cmd = Command::new("df");
	let output = cmd.output().await?;
    ensure!(output.status.success(), "df command failed");

    // sample
    // ファイルシス   1K-ブロック     使用    使用可 使用% マウント位置
    // /dev/root        122621412 12964620 104641120   12% /
    // devtmpfs           1800568        0   1800568    0% /dev
    // tmpfs              1965432        0   1965432    0% /dev/shm
    // tmpfs              1965432    17116   1948316    1% /run
    // tmpfs                 5120        4      5116    1% /run/lock
    // tmpfs              1965432        0   1965432    0% /sys/fs/cgroup
    // /dev/mmcblk0p1      258095    49324    208772   20% /boot
    // /dev/sda1         59280316 57109344         0  100% /media/usbbkup
    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut total = None;
    let mut avail = None;
    for line in stdout.lines().skip(1) {
        let mut total_tmp = None;
        let mut avail_tmp = None;
        let mut mp_tmp = None;
        for (col_no, elem) in line.split_ascii_whitespace().enumerate() {
            match col_no {
                1 => total_tmp = Some(elem),
                3 => avail_tmp = Some(elem),
                5 => mp_tmp = Some(elem),
                _ => (),
            }
        }
        if let Some(mp) = mp_tmp {
            if mp == "/" {
                total = total_tmp;
                avail = avail_tmp;
            }
        }
    }
    let total = total.ok_or(anyhow!("parse error"))?;
    let avail = avail.ok_or(anyhow!("parse error"))?;
    let total_gib = total.parse::<u64>()? as f64 / 1024.0 / 1024.0;
    let avail_gib = avail.parse::<u64>()? as f64 / 1024.0 / 1024.0;

    Ok(DiskInfo { total_gib, avail_gib })
}

#[derive(Debug, Clone, Copy)]
struct CpuTemp {
    temp: f64,
}

/// CPU 温度 (正確には違うかもしれない。ボード上の何らかの温度センサの値。) を取得する。
///
/// "/sys/class/thermal/thermal_zone0/temp" による。
/// デバイスファイルが存在しないと [std::io::Error] (NotFound) になる。
/// Linux 汎用のようだが少なくとも WSL2 では存在しない。
/// RasPi only で `vcgencmd measure_temp` という手もあるが、
/// 人が読みやすい代わりにパースが難しくなるのでデバイスファイルの方を使う。
async fn get_cpu_temp() -> Result<CpuTemp> {
    let buf = tokio::fs::read("/sys/class/thermal/thermal_zone0/temp").await?;
    let text = String::from_utf8_lossy(&buf);

    // 'C を 1000 倍した整数が得られるので変換する
    let temp: u64 = text.trim().parse()?;
    let temp = temp as f64 / 1000.0;

    Ok(CpuTemp { temp })
}
<commit_msg>氷点下の CPU 温度が返ると符号なし整数への変換エラーになるのを修正<commit_after>use crate::sys::config;
use crate::sys::taskserver::Control;
use super::SystemModule;
use anyhow::{Result, anyhow, ensure};
use chrono::NaiveTime;
use log::info;
use serde::{Serialize, Deserialize};
use tokio::process::Command;

#[derive(Clone, Serialize, Deserialize)]
struct HealthConfig {
    enabled: bool,
    debug_exec_once: bool,
}

pub struct Health {
    config: HealthConfig,
    wakeup_list: Vec<NaiveTime>,
}

impl Health {
    pub fn new(wakeup_list: Vec<NaiveTime>) -> Result<Self> {
        info!("[health] initialize");

        let jsobj = config::get_object(&["health"])
            .map_or(Err(anyhow!("Config not found: health")), Ok)?;
        let config: HealthConfig = serde_json::from_value(jsobj)?;

        Ok(Health {
            config,
            wakeup_list,
        })
    }

    async fn health_task(&mut self, ctrl: &Control) -> Result<()> {
        let cpu_info = get_cpu_info().await;
        info!("{:?}", cpu_info);
        let mem_info = get_mem_info().await;
        info!("{:?}", mem_info?);
        let disk_info = get_disk_info().await;
        info!("{:?}", disk_info?);
        let cpu_temp = get_cpu_temp().await;
        info!("{:?}", cpu_temp?);

        Ok(())
    }

    async fn health_task_entry(ctrl: Control) -> Result<()> {
        let mut health = ctrl.sysmods().health.write().await;
        health.health_task(&ctrl).await
    }
}

impl SystemModule for Health {
    fn on_start(&self, ctrl: &Control) {
        info!("[health] on_start");
        if self.config.enabled {
            if self.config.debug_exec_once {
                ctrl.spawn_oneshot_task(
                    "health_check",
                    Health::health_task_entry);
            }
            else {
                ctrl.spawn_periodic_task(
                    "health_check",
                    &self.wakeup_list,
                    Health::health_task_entry);
            }
        }
    }
}

#[derive(Debug, Clone)]
struct CpuInfo {
    cpu_percent_total: f64,
    cpu_percent_list: Vec<f64>,
}

async fn get_cpu_info() -> Result<CpuInfo> {
    let buf = tokio::fs::read("/proc/stat").await?;
    let text = String::from_utf8_lossy(&buf);

    // See `man proc`
    // user   (1) Time spent in user mode.
    // nice   (2) Time spent in user mode with low priority (nice).
    // system (3) Time spent in system mode.
    // idle   (4) Time spent in the idle task.  This value should be USER_HZ times the second entry in the /proc/uptime pseudo-file.
    // iowait (since Linux 2.5.41)
    //        (5) Time waiting for I/O to complete.  This value is not reliable, for the following reasons:
    // irq (since Linux 2.6.0-test4)
    //        (6) Time servicing interrupts.
    // softirq (since Linux 2.6.0-test4)
    //        (7) Time servicing softirqs.
    // steal (since Linux 2.6.11)
    //        (8)  Stolen  time,  which  is the time spent in other operating systems when running in a virtualized environment
    // guest (since Linux 2.6.24)
    //        (9) Time spent running a virtual CPU for guest operating systems under the control of the Linux kernel.
    // guest_nice (since Linux 2.6.33)
    //        (10)  Time spent running a niced guest (virtual CPU for guest operating systems under the
    //        control of the Linux kernel).
    let mut cpu_percent_total = None;
    let mut cpu_percent_list = vec![];
    for line in text.lines() {
        let mut name = None;
        let mut user = None;
        let mut nice = None;
        let mut system = None;
        let mut idle = None;
        for (col_no, elem) in line.split_ascii_whitespace().enumerate() {
            match col_no {
                0 => name = Some(elem),
                1 => user = Some(elem),
                2 => nice = Some(elem),
                3 => system = Some(elem),
                4 => idle = Some(elem),
                _ => (),
            }
        }
        // cpu or cpu%d の行を取り出す
        if name.is_none() || !name.unwrap().starts_with("cpu") {
            continue;
        }

        let user: u64 = user.ok_or(anyhow!("parse error"))?.parse()?;
        let nice: u64 = nice.ok_or(anyhow!("parse error"))?.parse()?;
        let system: u64 = system.ok_or(anyhow!("parse error"))?.parse()?;
        let idle: u64 = idle.ok_or(anyhow!("parse error"))?.parse()?;
        let total = user + nice + system + idle;
        let value = (total - idle) as f64 / total as f64;
        if name == Some("cpu") {
            cpu_percent_total = Some(value);
        }
        else {
            cpu_percent_list.push(value);
        }
    }

    ensure!(cpu_percent_total.is_some());
    ensure!(!cpu_percent_list.is_empty());
    let cpu_percent_total = cpu_percent_total.unwrap();
    Ok(CpuInfo {
        cpu_percent_total,
        cpu_percent_list,
    })
}

#[derive(Debug, Clone, Copy)]
struct MemInfo {
    total_mib: f64,
    avail_mib: f64,
}

async fn get_mem_info() -> Result<MemInfo> {
    let mut cmd = Command::new("free");
	let output = cmd.output().await?;
    ensure!(output.status.success(), "free command failed");

    // sample
    //                total        used        free      shared  buff/cache   available
    // Mem:        13034888     4119272     5561960          68     3353656     8609008
    // Swap:        4194304           0     4194304
    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut total = None;
    let mut avail = None;
    for (line_no, line) in stdout.lines().enumerate() {
        if line_no != 1 {
            continue;
        }
        for (col_no, elem) in line.split_ascii_whitespace().enumerate() {
            match col_no {
                1 => total = Some(elem),
                6 => avail = Some(elem),
                _ => (),
            }
        }
        break;
    }
    let total = total.ok_or(anyhow!("parse error"))?;
    let avail = avail.ok_or(anyhow!("parse error"))?;
    let total_mib = total.parse::<u64>()? as f64 / 1024.0;
    let avail_mib = avail.parse::<u64>()? as f64 / 1024.0;

    Ok(MemInfo { total_mib, avail_mib })
}

#[derive(Debug, Clone, Copy)]
struct DiskInfo {
    total_gib: f64,
    avail_gib: f64,
}

async fn get_disk_info() -> Result<DiskInfo> {
    let mut cmd = Command::new("df");
	let output = cmd.output().await?;
    ensure!(output.status.success(), "df command failed");

    // sample
    // ファイルシス   1K-ブロック     使用    使用可 使用% マウント位置
    // /dev/root        122621412 12964620 104641120   12% /
    // devtmpfs           1800568        0   1800568    0% /dev
    // tmpfs              1965432        0   1965432    0% /dev/shm
    // tmpfs              1965432    17116   1948316    1% /run
    // tmpfs                 5120        4      5116    1% /run/lock
    // tmpfs              1965432        0   1965432    0% /sys/fs/cgroup
    // /dev/mmcblk0p1      258095    49324    208772   20% /boot
    // /dev/sda1         59280316 57109344         0  100% /media/usbbkup
    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut total = None;
    let mut avail = None;
    for line in stdout.lines().skip(1) {
        let mut total_tmp = None;
        let mut avail_tmp = None;
        let mut mp_tmp = None;
        for (col_no, elem) in line.split_ascii_whitespace().enumerate() {
            match col_no {
                1 => total_tmp = Some(elem),
                3 => avail_tmp = Some(elem),
                5 => mp_tmp = Some(elem),
                _ => (),
            }
        }
        if let Some(mp) = mp_tmp {
            if mp == "/" {
                total = total_tmp;
                avail = avail_tmp;
            }
        }
    }
    let total = total.ok_or(anyhow!("parse error"))?;
    let avail = avail.ok_or(anyhow!("parse error"))?;
    let total_gib = total.parse::<u64>()? as f64 / 1024.0 / 1024.0;
    let avail_gib = avail.parse::<u64>()? as f64 / 1024.0 / 1024.0;

    Ok(DiskInfo { total_gib, avail_gib })
}

#[derive(Debug, Clone, Copy)]
struct CpuTemp {
    temp: f64,
}

/// CPU 温度 (正確には違うかもしれない。ボード上の何らかの温度センサの値。) を取得する。
///
/// "/sys/class/thermal/thermal_zone0/temp" による。
/// デバイスファイルが存在しないと [std::io::Error] (NotFound) になる。
/// Linux 汎用のようだが少なくとも WSL2 では存在しない。
/// RasPi only で `vcgencmd measure_temp` という手もあるが、
/// 人が読みやすい代わりにパースが難しくなるのでデバイスファイルの方を使う。
async fn get_cpu_temp() -> Result<CpuTemp> {
    let buf = tokio::fs::read("/sys/class/thermal/thermal_zone0/temp").await?;
    let text = String::from_utf8_lossy(&buf);

    // 'C を 1000 倍した整数が得られるので変換する
    let temp: f64 = text.trim().parse()?;
    let temp = temp as f64 / 1000.0;

    Ok(CpuTemp { temp })
}
<|endoftext|>


<commit_before>use super::Cli;
use libc::{kill, pid_t, SIGINT};
use std::io::{stderr, stdin, stdout, BufRead, BufReader, ErrorKind, Stdout, Write};
use std::process::Child;
use std::thread::spawn;
use termion::raw::{IntoRawMode, RawTerminal};

/// Reads the console messages from `child`'s standard output, sending SIGTERM
/// to the child when the process is terminated.
pub fn process(cli: &Cli, mut child: Child) {
    let raw_mode = forward_stdin_if_piped(&mut child);
    forward_stderr_if_piped(&mut child, raw_mode.is_some());
    let mut to_print = Vec::new();
    let mut reader = BufReader::new(child.stdout.as_mut().expect("Child's stdout not piped."));
    loop {
        let buffer = reader
            .fill_buf()
            .expect("Unable to read from child process.");
        if buffer.is_empty() {
            // The child process has closed its stdout, likely by exiting.
            break;
        }
        // Print the bytes received over stdout. If the terminal is in raw mode,
        // translate '\n' into '\r\n'.
        for &byte in buffer {
            if raw_mode.is_some() && byte == b'\n' {
                to_print.push(b'\r');
            }
            to_print.push(byte);
        }
        stdout()
            .write_all(&to_print)
            .expect("Unable to echo child's stdout.");
        to_print.clear();

        let buffer_len = buffer.len();
        reader.consume(buffer_len);
    }
    if cli.verbose {
        println!("Waiting for child process.\r");
    }
    let status = child.wait().expect("Unable to wait for child process");
    drop(raw_mode);
    assert!(
        status.success(),
        "Child process did not exit successfully. {}",
        status
    );
}

// If child's stdin is piped, this sets the terminal to raw mode and spawns a
// thread that forwards our stdin to child's stdin. The thread sends SIGINT to
// the child if Ctrl+C is pressed. Returns a RawTerminal, which reverts the
// terminal to its previous configuration on drop.
fn forward_stdin_if_piped(child: &mut Child) -> Option<RawTerminal<Stdout>> {
    let mut child_stdin = child.stdin.take()?;
    let child_id = child.id();
    spawn(move || {
        let our_stdin = stdin();
        let mut our_stdin = our_stdin.lock();
        loop {
            let buffer = our_stdin.fill_buf().expect("Failed to read stdin.");
            if buffer.is_empty() {
                // Our stdin was closed. We interpret this as a signal to exit,
                // because pressing Ctrl+C to trigger an exit is no longer
                // possible.
                break;
            }
            // In raw mode, pressing Ctrl+C will send a '3' byte to stdin ("end
            // of message" ASCII value) instead of sending SIGINT. Identify that
            // case, and exit if it occurs.
            if buffer.contains(&3) {
                break;
            }
            match child_stdin.write(buffer) {
                // A BrokenPipe error occurs when the child has exited. Exit
                // without sending SIGINT.
                Err(error) if error.kind() == ErrorKind::BrokenPipe => return,

                Err(error) => panic!("Failed to forward stdin: {}", error),
                Ok(bytes) => our_stdin.consume(bytes),
            }
        }
        // Send SIGINT to the child, telling it to exit. After the child exits,
        // the main loop will detect the exit and we will shut down cleanly.
        //
        // Safety: Sending SIGINT to a process is a safe operation -- kill is
        // marked unsafe because it is a FFI function.
        unsafe {
            kill(child_id as pid_t, SIGINT);
        }
    });
    Some(
        stdout()
            .into_raw_mode()
            .expect("Failed to set terminal to raw mode."),
    )
}

// Forwards child's stderr to our stderr if child's stderr is piped, converting
// line endings to CRLF if raw_mode is true.
fn forward_stderr_if_piped(child: &mut Child, raw_mode: bool) {
    let child_stderr = match child.stderr.take() {
        None => return,
        Some(child_stderr) => child_stderr,
    };
    spawn(move || {
        let mut to_print = Vec::new();
        let mut reader = BufReader::new(child_stderr);
        loop {
            let buffer = reader.fill_buf().expect("Unable to read child's stderr.");
            if buffer.is_empty() {
                return;
            }
            for &byte in buffer {
                if raw_mode && byte == b'\n' {
                    to_print.push(b'\r');
                }
                to_print.push(byte);
            }
            stderr()
                .write_all(&to_print)
                .expect("Unable to echo child's stderr.");
            to_print.clear();
            let buffer_len = buffer.len();
            reader.consume(buffer_len);
        }
    });
}
<commit_msg>Runner: flash stdout after each echo.<commit_after>use super::Cli;
use libc::{kill, pid_t, SIGINT};
use std::io::{stderr, stdin, stdout, BufRead, BufReader, ErrorKind, Stdout, Write};
use std::process::Child;
use std::thread::spawn;
use termion::raw::{IntoRawMode, RawTerminal};

/// Reads the console messages from `child`'s standard output, sending SIGTERM
/// to the child when the process is terminated.
pub fn process(cli: &Cli, mut child: Child) {
    let raw_mode = forward_stdin_if_piped(&mut child);
    forward_stderr_if_piped(&mut child, raw_mode.is_some());
    let mut to_print = Vec::new();
    let mut reader = BufReader::new(child.stdout.as_mut().expect("Child's stdout not piped."));
    loop {
        let buffer = reader
            .fill_buf()
            .expect("Unable to read from child process.");
        if buffer.is_empty() {
            // The child process has closed its stdout, likely by exiting.
            break;
        }
        // Print the bytes received over stdout. If the terminal is in raw mode,
        // translate '\n' into '\r\n'.
        for &byte in buffer {
            if raw_mode.is_some() && byte == b'\n' {
                to_print.push(b'\r');
            }
            to_print.push(byte);
        }
        let stdout = stdout();
        let mut lock = stdout.lock();
        lock.write_all(&to_print)
            .expect("Unable to echo child's stdout.");
        let _ = lock.flush();
        drop(lock);
        to_print.clear();

        let buffer_len = buffer.len();
        reader.consume(buffer_len);
    }
    if cli.verbose {
        println!("Waiting for child process.\r");
    }
    let status = child.wait().expect("Unable to wait for child process");
    drop(raw_mode);
    assert!(
        status.success(),
        "Child process did not exit successfully. {}",
        status
    );
}

// If child's stdin is piped, this sets the terminal to raw mode and spawns a
// thread that forwards our stdin to child's stdin. The thread sends SIGINT to
// the child if Ctrl+C is pressed. Returns a RawTerminal, which reverts the
// terminal to its previous configuration on drop.
fn forward_stdin_if_piped(child: &mut Child) -> Option<RawTerminal<Stdout>> {
    let mut child_stdin = child.stdin.take()?;
    let child_id = child.id();
    spawn(move || {
        let our_stdin = stdin();
        let mut our_stdin = our_stdin.lock();
        loop {
            let buffer = our_stdin.fill_buf().expect("Failed to read stdin.");
            if buffer.is_empty() {
                // Our stdin was closed. We interpret this as a signal to exit,
                // because pressing Ctrl+C to trigger an exit is no longer
                // possible.
                break;
            }
            // In raw mode, pressing Ctrl+C will send a '3' byte to stdin ("end
            // of message" ASCII value) instead of sending SIGINT. Identify that
            // case, and exit if it occurs.
            if buffer.contains(&3) {
                break;
            }
            match child_stdin.write(buffer) {
                // A BrokenPipe error occurs when the child has exited. Exit
                // without sending SIGINT.
                Err(error) if error.kind() == ErrorKind::BrokenPipe => return,

                Err(error) => panic!("Failed to forward stdin: {}", error),
                Ok(bytes) => our_stdin.consume(bytes),
            }
        }
        // Send SIGINT to the child, telling it to exit. After the child exits,
        // the main loop will detect the exit and we will shut down cleanly.
        //
        // Safety: Sending SIGINT to a process is a safe operation -- kill is
        // marked unsafe because it is a FFI function.
        unsafe {
            kill(child_id as pid_t, SIGINT);
        }
    });
    Some(
        stdout()
            .into_raw_mode()
            .expect("Failed to set terminal to raw mode."),
    )
}

// Forwards child's stderr to our stderr if child's stderr is piped, converting
// line endings to CRLF if raw_mode is true.
fn forward_stderr_if_piped(child: &mut Child, raw_mode: bool) {
    let child_stderr = match child.stderr.take() {
        None => return,
        Some(child_stderr) => child_stderr,
    };
    spawn(move || {
        let mut to_print = Vec::new();
        let mut reader = BufReader::new(child_stderr);
        loop {
            let buffer = reader.fill_buf().expect("Unable to read child's stderr.");
            if buffer.is_empty() {
                return;
            }
            for &byte in buffer {
                if raw_mode && byte == b'\n' {
                    to_print.push(b'\r');
                }
                to_print.push(byte);
            }
            stderr()
                .write_all(&to_print)
                .expect("Unable to echo child's stderr.");
            to_print.clear();
            let buffer_len = buffer.len();
            reader.consume(buffer_len);
        }
    });
}
<|endoftext|>


<commit_before>// Copyright 2020 The Tink-Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#![feature(test)]
extern crate test;
use test::Bencher;

const MSG: &[u8] = b"this data needs to be authenticated";

fn setup(kt: tink_proto::KeyTemplate) -> (Box<dyn tink_core::Mac>, Vec<u8>) {
    tink_mac::init();
    let kh = tink_core::keyset::Handle::new(&kt).unwrap();
    let m = tink_mac::new(&kh).unwrap();

    let tag = m.compute_mac(MSG).unwrap();
    (m, tag)
}

#[bench]
fn bench_hmac_sha256_tag128_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha256_tag128_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha256_tag128_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha256_tag128_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_hmac_sha256_tag256_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha256_tag256_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha256_tag256_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha256_tag256_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag256_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha512_tag256_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag256_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha512_tag256_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag512_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha512_tag512_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag512_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha512_tag512_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_aes_cmac_tag128_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::aes_cmac_tag128_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_aes_cmac_tag128_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::aes_cmac_tag128_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}
<commit_msg>tink-mac: add benchmark for verify failure<commit_after>// Copyright 2020 The Tink-Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#![feature(test)]
extern crate test;
use test::Bencher;

const MSG: &[u8] = b"this data needs to be authenticated";

fn setup(kt: tink_proto::KeyTemplate) -> (Box<dyn tink_core::Mac>, Vec<u8>) {
    tink_mac::init();
    let kh = tink_core::keyset::Handle::new(&kt).unwrap();
    let m = tink_mac::new(&kh).unwrap();

    let tag = m.compute_mac(MSG).unwrap();
    (m, tag)
}

/// Size of the prefix information in the tag. If this is corrupted, the tag will be
/// rejected immediately without performing any cryptographic operations.
const PREFIX_SIZE: usize = tink_core::cryptofmt::NON_RAW_PREFIX_SIZE;

fn setup_failure(kt: tink_proto::KeyTemplate) -> (Box<dyn tink_core::Mac>, Vec<u8>) {
    let (m, tag) = setup(kt);
    (
        m,
        tag.iter()
            .enumerate()
            .map(|(i, b)| if i < PREFIX_SIZE { *b } else { b ^ 0b10101010 })
            .collect(),
    )
}

#[bench]
fn bench_hmac_sha256_tag128_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha256_tag128_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha256_tag128_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha256_tag128_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_hmac_sha256_tag128_mac_verify_fail(b: &mut Bencher) {
    let (m, tag) = setup_failure(tink_mac::hmac_sha256_tag128_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap_err());
}

#[bench]
fn bench_hmac_sha256_tag256_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha256_tag256_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha256_tag256_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha256_tag256_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_hmac_sha256_tag256_mac_verify_fail(b: &mut Bencher) {
    let (m, tag) = setup_failure(tink_mac::hmac_sha256_tag256_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap_err());
}

#[bench]
fn bench_hmac_sha512_tag256_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha512_tag256_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag256_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha512_tag256_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag256_mac_verify_fail(b: &mut Bencher) {
    let (m, tag) = setup_failure(tink_mac::hmac_sha512_tag256_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap_err());
}

#[bench]
fn bench_hmac_sha512_tag512_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::hmac_sha512_tag512_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag512_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::hmac_sha512_tag512_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_hmac_sha512_tag512_mac_verify_fail(b: &mut Bencher) {
    let (m, tag) = setup_failure(tink_mac::hmac_sha512_tag512_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap_err());
}

#[bench]
fn bench_aes_cmac_tag128_mac_compute(b: &mut Bencher) {
    let (m, _tag) = setup(tink_mac::aes_cmac_tag128_key_template());
    b.iter(|| m.compute_mac(MSG).unwrap());
}

#[bench]
fn bench_aes_cmac_tag128_mac_verify(b: &mut Bencher) {
    let (m, tag) = setup(tink_mac::aes_cmac_tag128_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap());
}

#[bench]
fn bench_aes_cmac_tag128_mac_verify_fail(b: &mut Bencher) {
    let (m, tag) = setup_failure(tink_mac::aes_cmac_tag128_key_template());
    b.iter(|| m.verify_mac(&tag, MSG).unwrap_err());
}
<|endoftext|>


<commit_before>//! Workspace information we get from cargo consists of two pieces. The first is
//! the output of `cargo metadata`. The second is the output of running
//! `build.rs` files (`OUT_DIR` env var, extra cfg flags) and compiling proc
//! macro.
//!
//! This module implements this second part. We use "build script" terminology
//! here, but it covers procedural macros as well.

use std::{
    cell::RefCell,
    io,
    path::PathBuf,
    process::{Command, Stdio},
};

use cargo_metadata::{camino::Utf8Path, Message};
use la_arena::ArenaMap;
use paths::AbsPathBuf;
use rustc_hash::FxHashMap;
use serde::Deserialize;

use crate::{cfg_flag::CfgFlag, CargoConfig, CargoWorkspace, Package};

#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct WorkspaceBuildScripts {
    pub(crate) outputs: ArenaMap<Package, BuildScriptOutput>,
    error: Option<String>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct BuildScriptOutput {
    /// List of config flags defined by this package's build script.
    pub(crate) cfgs: Vec<CfgFlag>,
    /// List of cargo-related environment variables with their value.
    ///
    /// If the package has a build script which defines environment variables,
    /// they can also be found here.
    pub(crate) envs: Vec<(String, String)>,
    /// Directory where a build script might place its output.
    pub(crate) out_dir: Option<AbsPathBuf>,
    /// Path to the proc-macro library file if this package exposes proc-macros.
    pub(crate) proc_macro_dylib_path: Option<AbsPathBuf>,
}

impl WorkspaceBuildScripts {
    fn build_command(config: &CargoConfig) -> Command {
        if let Some([program, args @ ..]) = config.run_build_script_command.as_deref() {
            let mut cmd = Command::new(program);
            cmd.args(args);
            return cmd;
        }

        let mut cmd = Command::new(toolchain::cargo());

        cmd.args(&["check", "--quiet", "--workspace", "--message-format=json"]);

        // --all-targets includes tests, benches and examples in addition to the
        // default lib and bins. This is an independent concept from the --targets
        // flag below.
        cmd.arg("--all-targets");

        if let Some(target) = &config.target {
            cmd.args(&["--target", target]);
        }

        if config.all_features {
            cmd.arg("--all-features");
        } else {
            if config.no_default_features {
                cmd.arg("--no-default-features");
            }
            if !config.features.is_empty() {
                cmd.arg("--features");
                cmd.arg(config.features.join(" "));
            }
        }

        cmd
    }
    pub(crate) fn run(
        config: &CargoConfig,
        workspace: &CargoWorkspace,
        progress: &dyn Fn(String),
    ) -> io::Result<WorkspaceBuildScripts> {
        let mut cmd = Self::build_command(config);

        if config.wrap_rustc_in_build_scripts {
            // Setup RUSTC_WRAPPER to point to `rust-analyzer` binary itself. We use
            // that to compile only proc macros and build scripts during the initial
            // `cargo check`.
            let myself = std::env::current_exe()?;
            cmd.env("RUSTC_WRAPPER", myself);
            cmd.env("RA_RUSTC_WRAPPER", "1");
        }

        cmd.current_dir(workspace.workspace_root());

        cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).stdin(Stdio::null());

        let mut res = WorkspaceBuildScripts::default();
        // NB: Cargo.toml could have been modified between `cargo metadata` and
        // `cargo check`. We shouldn't assume that package ids we see here are
        // exactly those from `config`.
        let mut by_id: FxHashMap<String, Package> = FxHashMap::default();

        for package in workspace.packages() {
            res.outputs.insert(package, BuildScriptOutput::default());
            by_id.insert(workspace[package].id.clone(), package);
        }

        let errors = RefCell::new(String::new());
        let push_err = |err: &str| {
            let mut e = errors.borrow_mut();
            e.push_str(err);
            e.push('\n');
        };
        let output = stdx::process::streaming_output(
            cmd,
            &mut |line| {
                // Copy-pasted from existing cargo_metadata. It seems like we
                // should be using serde_stacker here?
                let mut deserializer = serde_json::Deserializer::from_str(line);
                deserializer.disable_recursion_limit();
                let message = Message::deserialize(&mut deserializer)
                    .unwrap_or_else(|_| Message::TextLine(line.to_string()));

                match message {
                    Message::BuildScriptExecuted(message) => {
                        let package = match by_id.get(&message.package_id.repr) {
                            Some(&it) => it,
                            None => return,
                        };
                        let cfgs = {
                            let mut acc = Vec::new();
                            for cfg in message.cfgs {
                                match cfg.parse::<CfgFlag>() {
                                    Ok(it) => acc.push(it),
                                    Err(err) => {
                                        push_err(&format!(
                                            "invalid cfg from cargo-metadata: {}",
                                            err
                                        ));
                                        return;
                                    }
                                };
                            }
                            acc
                        };
                        let package_build_data = &mut res.outputs[package];
                        // cargo_metadata crate returns default (empty) path for
                        // older cargos, which is not absolute, so work around that.
                        if !message.out_dir.as_str().is_empty() {
                            let out_dir =
                                AbsPathBuf::assert(PathBuf::from(message.out_dir.into_os_string()));
                            package_build_data.out_dir = Some(out_dir);
                            package_build_data.cfgs = cfgs;
                        }

                        package_build_data.envs = message.env;
                    }
                    Message::CompilerArtifact(message) => {
                        let package = match by_id.get(&message.package_id.repr) {
                            Some(it) => *it,
                            None => return,
                        };

                        progress(format!("metadata {}", message.target.name));

                        if message.target.kind.iter().any(|k| k == "proc-macro") {
                            // Skip rmeta file
                            if let Some(filename) =
                                message.filenames.iter().find(|name| is_dylib(name))
                            {
                                let filename = AbsPathBuf::assert(PathBuf::from(&filename));
                                res.outputs[package].proc_macro_dylib_path = Some(filename);
                            }
                        }
                    }
                    Message::CompilerMessage(message) => {
                        progress(message.target.name);

                        if let Some(diag) = message.message.rendered.as_deref() {
                            push_err(diag);
                        }
                    }
                    Message::BuildFinished(_) => {}
                    Message::TextLine(_) => {}
                    _ => {}
                }
            },
            &mut |line| {
                push_err(line);
            },
        )?;

        for package in workspace.packages() {
            let package_build_data = &mut res.outputs[package];
            tracing::info!(
                "{} BuildScriptOutput: {:?}",
                workspace[package].manifest.parent().display(),
                package_build_data,
            );
            // inject_cargo_env(package, package_build_data);
            if let Some(out_dir) = &package_build_data.out_dir {
                // NOTE: cargo and rustc seem to hide non-UTF-8 strings from env! and option_env!()
                if let Some(out_dir) = out_dir.as_os_str().to_str().map(|s| s.to_owned()) {
                    package_build_data.envs.push(("OUT_DIR".to_string(), out_dir));
                }
            }
        }

        let mut errors = errors.into_inner();
        if !output.status.success() {
            if errors.is_empty() {
                errors = "cargo check failed".to_string();
            }
            res.error = Some(errors);
        }

        Ok(res)
    }

    pub fn error(&self) -> Option<&str> {
        self.error.as_deref()
    }
}

// FIXME: File a better way to know if it is a dylib.
fn is_dylib(path: &Utf8Path) -> bool {
    match path.extension().map(|e| e.to_string().to_lowercase()) {
        None => false,
        Some(ext) => matches!(ext.as_str(), "dll" | "dylib" | "so"),
    }
}
<commit_msg>minor: log build script invocations<commit_after>//! Workspace information we get from cargo consists of two pieces. The first is
//! the output of `cargo metadata`. The second is the output of running
//! `build.rs` files (`OUT_DIR` env var, extra cfg flags) and compiling proc
//! macro.
//!
//! This module implements this second part. We use "build script" terminology
//! here, but it covers procedural macros as well.

use std::{cell::RefCell, io, path::PathBuf, process::Command};

use cargo_metadata::{camino::Utf8Path, Message};
use la_arena::ArenaMap;
use paths::AbsPathBuf;
use rustc_hash::FxHashMap;
use serde::Deserialize;

use crate::{cfg_flag::CfgFlag, CargoConfig, CargoWorkspace, Package};

#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct WorkspaceBuildScripts {
    pub(crate) outputs: ArenaMap<Package, BuildScriptOutput>,
    error: Option<String>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct BuildScriptOutput {
    /// List of config flags defined by this package's build script.
    pub(crate) cfgs: Vec<CfgFlag>,
    /// List of cargo-related environment variables with their value.
    ///
    /// If the package has a build script which defines environment variables,
    /// they can also be found here.
    pub(crate) envs: Vec<(String, String)>,
    /// Directory where a build script might place its output.
    pub(crate) out_dir: Option<AbsPathBuf>,
    /// Path to the proc-macro library file if this package exposes proc-macros.
    pub(crate) proc_macro_dylib_path: Option<AbsPathBuf>,
}

impl WorkspaceBuildScripts {
    fn build_command(config: &CargoConfig) -> Command {
        if let Some([program, args @ ..]) = config.run_build_script_command.as_deref() {
            let mut cmd = Command::new(program);
            cmd.args(args);
            return cmd;
        }

        let mut cmd = Command::new(toolchain::cargo());

        cmd.args(&["check", "--quiet", "--workspace", "--message-format=json"]);

        // --all-targets includes tests, benches and examples in addition to the
        // default lib and bins. This is an independent concept from the --targets
        // flag below.
        cmd.arg("--all-targets");

        if let Some(target) = &config.target {
            cmd.args(&["--target", target]);
        }

        if config.all_features {
            cmd.arg("--all-features");
        } else {
            if config.no_default_features {
                cmd.arg("--no-default-features");
            }
            if !config.features.is_empty() {
                cmd.arg("--features");
                cmd.arg(config.features.join(" "));
            }
        }

        cmd
    }
    pub(crate) fn run(
        config: &CargoConfig,
        workspace: &CargoWorkspace,
        progress: &dyn Fn(String),
    ) -> io::Result<WorkspaceBuildScripts> {
        let mut cmd = Self::build_command(config);

        if config.wrap_rustc_in_build_scripts {
            // Setup RUSTC_WRAPPER to point to `rust-analyzer` binary itself. We use
            // that to compile only proc macros and build scripts during the initial
            // `cargo check`.
            let myself = std::env::current_exe()?;
            cmd.env("RUSTC_WRAPPER", myself);
            cmd.env("RA_RUSTC_WRAPPER", "1");
        }

        cmd.current_dir(workspace.workspace_root());

        let mut res = WorkspaceBuildScripts::default();
        // NB: Cargo.toml could have been modified between `cargo metadata` and
        // `cargo check`. We shouldn't assume that package ids we see here are
        // exactly those from `config`.
        let mut by_id: FxHashMap<String, Package> = FxHashMap::default();

        for package in workspace.packages() {
            res.outputs.insert(package, BuildScriptOutput::default());
            by_id.insert(workspace[package].id.clone(), package);
        }

        let errors = RefCell::new(String::new());
        let push_err = |err: &str| {
            let mut e = errors.borrow_mut();
            e.push_str(err);
            e.push('\n');
        };

        tracing::info!("Running build scripts: {:?}", cmd);
        let output = stdx::process::streaming_output(
            cmd,
            &mut |line| {
                // Copy-pasted from existing cargo_metadata. It seems like we
                // should be using serde_stacker here?
                let mut deserializer = serde_json::Deserializer::from_str(line);
                deserializer.disable_recursion_limit();
                let message = Message::deserialize(&mut deserializer)
                    .unwrap_or_else(|_| Message::TextLine(line.to_string()));

                match message {
                    Message::BuildScriptExecuted(message) => {
                        let package = match by_id.get(&message.package_id.repr) {
                            Some(&it) => it,
                            None => return,
                        };
                        let cfgs = {
                            let mut acc = Vec::new();
                            for cfg in message.cfgs {
                                match cfg.parse::<CfgFlag>() {
                                    Ok(it) => acc.push(it),
                                    Err(err) => {
                                        push_err(&format!(
                                            "invalid cfg from cargo-metadata: {}",
                                            err
                                        ));
                                        return;
                                    }
                                };
                            }
                            acc
                        };
                        let package_build_data = &mut res.outputs[package];
                        // cargo_metadata crate returns default (empty) path for
                        // older cargos, which is not absolute, so work around that.
                        if !message.out_dir.as_str().is_empty() {
                            let out_dir =
                                AbsPathBuf::assert(PathBuf::from(message.out_dir.into_os_string()));
                            package_build_data.out_dir = Some(out_dir);
                            package_build_data.cfgs = cfgs;
                        }

                        package_build_data.envs = message.env;
                    }
                    Message::CompilerArtifact(message) => {
                        let package = match by_id.get(&message.package_id.repr) {
                            Some(it) => *it,
                            None => return,
                        };

                        progress(format!("metadata {}", message.target.name));

                        if message.target.kind.iter().any(|k| k == "proc-macro") {
                            // Skip rmeta file
                            if let Some(filename) =
                                message.filenames.iter().find(|name| is_dylib(name))
                            {
                                let filename = AbsPathBuf::assert(PathBuf::from(&filename));
                                res.outputs[package].proc_macro_dylib_path = Some(filename);
                            }
                        }
                    }
                    Message::CompilerMessage(message) => {
                        progress(message.target.name);

                        if let Some(diag) = message.message.rendered.as_deref() {
                            push_err(diag);
                        }
                    }
                    Message::BuildFinished(_) => {}
                    Message::TextLine(_) => {}
                    _ => {}
                }
            },
            &mut |line| {
                push_err(line);
            },
        )?;

        for package in workspace.packages() {
            let package_build_data = &mut res.outputs[package];
            tracing::info!(
                "{} BuildScriptOutput: {:?}",
                workspace[package].manifest.parent().display(),
                package_build_data,
            );
            // inject_cargo_env(package, package_build_data);
            if let Some(out_dir) = &package_build_data.out_dir {
                // NOTE: cargo and rustc seem to hide non-UTF-8 strings from env! and option_env!()
                if let Some(out_dir) = out_dir.as_os_str().to_str().map(|s| s.to_owned()) {
                    package_build_data.envs.push(("OUT_DIR".to_string(), out_dir));
                }
            }
        }

        let mut errors = errors.into_inner();
        if !output.status.success() {
            if errors.is_empty() {
                errors = "cargo check failed".to_string();
            }
            res.error = Some(errors);
        }

        Ok(res)
    }

    pub fn error(&self) -> Option<&str> {
        self.error.as_deref()
    }
}

// FIXME: File a better way to know if it is a dylib.
fn is_dylib(path: &Utf8Path) -> bool {
    match path.extension().map(|e| e.to_string().to_lowercase()) {
        None => false,
        Some(ext) => matches!(ext.as_str(), "dll" | "dylib" | "so"),
    }
}
<|endoftext|>


<commit_before>//! Builtin macro
use crate::db::AstDatabase;
use crate::{
    ast::{self, AstNode},
    name, AstId, CrateId, HirFileId, MacroCallId, MacroDefId, MacroDefKind, MacroFileKind,
    TextUnit,
};

use crate::quote;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BuiltinFnLikeExpander {
    Column,
    File,
    Line,
    Stringify,
}

struct BuiltInMacroInfo {
    name: name::Name,
    kind: BuiltinFnLikeExpander,
    expand: fn(
        db: &dyn AstDatabase,
        id: MacroCallId,
        _tt: &tt::Subtree,
    ) -> Result<tt::Subtree, mbe::ExpandError>,
}

macro_rules! register_builtin {
    ( $(($name:ident, $kind: ident) => $expand:ident),* ) => {
        const BUILTIN_MACROS: &[BuiltInMacroInfo] = &[
            $(BuiltInMacroInfo { name: name::$name, kind: BuiltinFnLikeExpander::$kind, expand: $expand }),*
        ];
    };
}

register_builtin! {
    (COLUMN_MACRO, Column) => column_expand,
    (FILE_MACRO, File) => file_expand,
    (LINE_MACRO, Line) => line_expand,
    (STRINGIFY_MACRO, Stringify) => stringify_expand
}

impl BuiltinFnLikeExpander {
    pub fn expand(
        &self,
        db: &dyn AstDatabase,
        id: MacroCallId,
        tt: &tt::Subtree,
    ) -> Result<tt::Subtree, mbe::ExpandError> {
        let expander = BUILTIN_MACROS
            .iter()
            .find(|it| *self == it.kind)
            .map(|it| it.expand)
            .ok_or_else(|| mbe::ExpandError::ConversionError)?;

        expander(db, id, tt)
    }
}

pub fn find_builtin_macro(
    ident: &name::Name,
    krate: CrateId,
    ast_id: AstId<ast::MacroCall>,
) -> Option<MacroDefId> {
    let kind = BUILTIN_MACROS.iter().find(|it| *ident == it.name).map(|it| it.kind)?;

    Some(MacroDefId { krate, ast_id, kind: MacroDefKind::BuiltIn(kind) })
}

fn to_line_number(db: &dyn AstDatabase, file: HirFileId, pos: TextUnit) -> usize {
    // FIXME: Use expansion info
    let file_id = file.original_file(db);
    let text = db.file_text(file_id);
    let mut line_num = 1;

    // Count line end
    for (i, c) in text.chars().enumerate() {
        if i == pos.to_usize() {
            break;
        }
        if c == '\n' {
            line_num += 1;
        }
    }

    line_num
}

fn line_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let arg = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;
    let arg_start = arg.syntax().text_range().start();

    let file = id.as_file(MacroFileKind::Expr);
    let line_num = to_line_number(db, file, arg_start);

    let expanded = quote! {
        #line_num
    };

    Ok(expanded)
}

fn stringify_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let macro_content = {
        let arg = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;
        let macro_args = arg.syntax().clone();
        let text = macro_args.text();
        let without_parens = TextUnit::of_char('(')..text.len() - TextUnit::of_char(')');
        text.slice(without_parens).to_string()
    };

    let expanded = quote! {
        #macro_content
    };

    Ok(expanded)
}

fn to_col_number(db: &dyn AstDatabase, file: HirFileId, pos: TextUnit) -> usize {
    // FIXME: Use expansion info
    let file_id = file.original_file(db);
    let text = db.file_text(file_id);
    let mut col_num = 1;

    for c in text[..pos.to_usize()].chars().rev() {
        if c == '\n' {
            break;
        }
        col_num = col_num + 1;
    }

    col_num
}

fn column_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let _arg = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;
    let col_start = macro_call.syntax().text_range().start();

    let file = id.as_file(MacroFileKind::Expr);
    let col_num = to_col_number(db, file, col_start);

    let expanded = quote! {
        #col_num
    };

    Ok(expanded)
}

fn file_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let _ = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;

    // FIXME: RA purposefully lacks knowledge of absolute file names
    // so just return "".
    let file_name = "";

    let expanded = quote! {
        #file_name
    };

    Ok(expanded)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{test_db::TestDB, MacroCallLoc};
    use ra_db::{fixture::WithFixture, SourceDatabase};

    fn expand_builtin_macro(s: &str, expander: BuiltinFnLikeExpander) -> String {
        let (db, file_id) = TestDB::with_single_file(&s);
        let parsed = db.parse(file_id);
        let macro_calls: Vec<_> =
            parsed.syntax_node().descendants().filter_map(|it| ast::MacroCall::cast(it)).collect();

        let ast_id_map = db.ast_id_map(file_id.into());

        // the first one should be a macro_rules
        let def = MacroDefId {
            krate: CrateId(0),
            ast_id: AstId::new(file_id.into(), ast_id_map.ast_id(&macro_calls[0])),
            kind: MacroDefKind::BuiltIn(expander),
        };

        let loc = MacroCallLoc {
            def,
            ast_id: AstId::new(file_id.into(), ast_id_map.ast_id(&macro_calls[1])),
        };

        let id = db.intern_macro(loc);
        let parsed = db.parse_or_expand(id.as_file(MacroFileKind::Expr)).unwrap();

        parsed.text().to_string()
    }

    #[test]
    fn test_column_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! column {() => {}}
        column!()
"#,
            BuiltinFnLikeExpander::Column,
        );

        assert_eq!(expanded, "9");
    }

    #[test]
    fn test_line_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! line {() => {}}
        line!()
"#,
            BuiltinFnLikeExpander::Line,
        );

        assert_eq!(expanded, "4");
    }

    #[test]
    fn test_stringify_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! stringify {() => {}}
        stringify!(a b c)
"#,
            BuiltinFnLikeExpander::Stringify,
        );

        assert_eq!(expanded, "\"a b c\"");
    }

    #[test]
    fn test_file_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! file {() => {}}
        file!()
"#,
            BuiltinFnLikeExpander::File,
        );

        assert_eq!(expanded, "\"\"");
    }
}
<commit_msg>Use macro for all the things<commit_after>//! Builtin macro
use crate::db::AstDatabase;
use crate::{
    ast::{self, AstNode},
    name, AstId, CrateId, HirFileId, MacroCallId, MacroDefId, MacroDefKind, MacroFileKind,
    TextUnit,
};

use crate::quote;

macro_rules! register_builtin {
    ( $(($name:ident, $kind: ident) => $expand:ident),* ) => {
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
        pub enum BuiltinFnLikeExpander {
            $($kind),*
        }

        impl BuiltinFnLikeExpander {
            pub fn expand(
                &self,
                db: &dyn AstDatabase,
                id: MacroCallId,
                tt: &tt::Subtree,
            ) -> Result<tt::Subtree, mbe::ExpandError> {
                let expander = match *self {
                    $( BuiltinFnLikeExpander::$kind => $expand, )*
                };
                expander(db, id, tt)
            }
        }

        pub fn find_builtin_macro(
            ident: &name::Name,
            krate: CrateId,
            ast_id: AstId<ast::MacroCall>,
        ) -> Option<MacroDefId> {
            let kind = match ident {
                 $( id if id == &name::$name => BuiltinFnLikeExpander::$kind, )*
                 _ => return None,
            };

            Some(MacroDefId { krate, ast_id, kind: MacroDefKind::BuiltIn(kind) })
        }
    };
}

register_builtin! {
    (COLUMN_MACRO, Column) => column_expand,
    (FILE_MACRO, File) => file_expand,
    (LINE_MACRO, Line) => line_expand,
    (STRINGIFY_MACRO, Stringify) => stringify_expand
}

fn to_line_number(db: &dyn AstDatabase, file: HirFileId, pos: TextUnit) -> usize {
    // FIXME: Use expansion info
    let file_id = file.original_file(db);
    let text = db.file_text(file_id);
    let mut line_num = 1;

    // Count line end
    for (i, c) in text.chars().enumerate() {
        if i == pos.to_usize() {
            break;
        }
        if c == '\n' {
            line_num += 1;
        }
    }

    line_num
}

fn line_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let arg = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;
    let arg_start = arg.syntax().text_range().start();

    let file = id.as_file(MacroFileKind::Expr);
    let line_num = to_line_number(db, file, arg_start);

    let expanded = quote! {
        #line_num
    };

    Ok(expanded)
}

fn stringify_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let macro_content = {
        let arg = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;
        let macro_args = arg.syntax().clone();
        let text = macro_args.text();
        let without_parens = TextUnit::of_char('(')..text.len() - TextUnit::of_char(')');
        text.slice(without_parens).to_string()
    };

    let expanded = quote! {
        #macro_content
    };

    Ok(expanded)
}

fn to_col_number(db: &dyn AstDatabase, file: HirFileId, pos: TextUnit) -> usize {
    // FIXME: Use expansion info
    let file_id = file.original_file(db);
    let text = db.file_text(file_id);
    let mut col_num = 1;

    for c in text[..pos.to_usize()].chars().rev() {
        if c == '\n' {
            break;
        }
        col_num = col_num + 1;
    }

    col_num
}

fn column_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let _arg = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;
    let col_start = macro_call.syntax().text_range().start();

    let file = id.as_file(MacroFileKind::Expr);
    let col_num = to_col_number(db, file, col_start);

    let expanded = quote! {
        #col_num
    };

    Ok(expanded)
}

fn file_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
    _tt: &tt::Subtree,
) -> Result<tt::Subtree, mbe::ExpandError> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);

    let _ = macro_call.token_tree().ok_or_else(|| mbe::ExpandError::UnexpectedToken)?;

    // FIXME: RA purposefully lacks knowledge of absolute file names
    // so just return "".
    let file_name = "";

    let expanded = quote! {
        #file_name
    };

    Ok(expanded)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{test_db::TestDB, MacroCallLoc};
    use ra_db::{fixture::WithFixture, SourceDatabase};

    fn expand_builtin_macro(s: &str, expander: BuiltinFnLikeExpander) -> String {
        let (db, file_id) = TestDB::with_single_file(&s);
        let parsed = db.parse(file_id);
        let macro_calls: Vec<_> =
            parsed.syntax_node().descendants().filter_map(|it| ast::MacroCall::cast(it)).collect();

        let ast_id_map = db.ast_id_map(file_id.into());

        // the first one should be a macro_rules
        let def = MacroDefId {
            krate: CrateId(0),
            ast_id: AstId::new(file_id.into(), ast_id_map.ast_id(&macro_calls[0])),
            kind: MacroDefKind::BuiltIn(expander),
        };

        let loc = MacroCallLoc {
            def,
            ast_id: AstId::new(file_id.into(), ast_id_map.ast_id(&macro_calls[1])),
        };

        let id = db.intern_macro(loc);
        let parsed = db.parse_or_expand(id.as_file(MacroFileKind::Expr)).unwrap();

        parsed.text().to_string()
    }

    #[test]
    fn test_column_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! column {() => {}}
        column!()
"#,
            BuiltinFnLikeExpander::Column,
        );

        assert_eq!(expanded, "9");
    }

    #[test]
    fn test_line_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! line {() => {}}
        line!()
"#,
            BuiltinFnLikeExpander::Line,
        );

        assert_eq!(expanded, "4");
    }

    #[test]
    fn test_stringify_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! stringify {() => {}}
        stringify!(a b c)
"#,
            BuiltinFnLikeExpander::Stringify,
        );

        assert_eq!(expanded, "\"a b c\"");
    }

    #[test]
    fn test_file_expand() {
        let expanded = expand_builtin_macro(
            r#"
        #[rustc_builtin_macro]
        macro_rules! file {() => {}}
        file!()
"#,
            BuiltinFnLikeExpander::File,
        );

        assert_eq!(expanded, "\"\"");
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{BufReader};
use std::path::PathBuf;

use cgmath::{Matrix4, Deg, Vector4, SquareMatrix, Vector3, Euler, Quaternion};
use soto::task::{task_log};
use soto::Error;
use sotolib_fbx::{RawFbx, id_name, friendly_name, ObjectTreeNode};
use sotolib_fbx::animation::{Animation};
use sotolib_fbx::simple::{Object, SimpleFbx, ObjectType, ModelProperties, Geometry};
use sotolib_smd::{Smd, SmdVertex, SmdTriangle, SmdAnimationFrameBone, SmdBone};

pub fn create_reference_smd(fbx: &PathBuf) -> Result<Smd, Error> {
    // Read in the fbx we got told to convert
    let file = BufReader::new(File::open(&fbx).unwrap());
    let fbx = SimpleFbx::from_raw(&RawFbx::parse(file).unwrap()).unwrap();
    let fbx_tree = ObjectTreeNode::from_simple(&fbx);

    // Go over all FBX root nodes and turn them into SMD data
    let mut smd = Smd::new();
    process_fbx_node(
        &fbx,
        &fbx_tree, &mut smd,
        &Matrix4::identity(),
        None
    )?;

    Ok(smd)
}

pub fn create_animation_smd(ref_smd: &Smd, fbx: &PathBuf) -> Result<Smd, Error> {
    // Read in the fbx we got told to convert
    let file = BufReader::new(File::open(&fbx).unwrap());
    let mut fbx = SimpleFbx::from_raw(&RawFbx::parse(file).unwrap()).unwrap();

    // Read in the animation data itself
    let animation = Animation::from_simple(&fbx).unwrap();

    // Count and log frames
    let frame_count = animation.frame_count(&fbx);
    task_log(format!("Animation has {} frames", frame_count));

    // Copy over every bone to the new animation SMD
    let mut smd = Smd::new();
    for bone in &ref_smd.bones {
        smd.bones.push(bone.clone());
    }

    // Finally, turn the animation data into bone positions in the SMD
    for frame in 0..frame_count {
        // First transform the FBX for this frame
        animation.transform_fbx_to_frame(&mut fbx, frame);

        // Now go over all models
        for (_, model) in fbx.objects.iter().filter(|&(_, o)| o.class.type_name() == "Model") {
            // For this model, look up the matching BoneId in the reference SMD
            if let Some(bone_id) = ref_smd.id_of_bone(&id_name(&model.name).unwrap()) {
                // Now that we have a model and a bone, we need the current translation and rotation
                // for the model
                let (translation, rotation) = calculate_animation_transforms_for(model);

                // And now that we have those, finally add the bone data to the animation SMD
                smd.set_animation(frame, bone_id, SmdAnimationFrameBone {
                    translation: translation.into(),
                    rotation: rotation.into(),
                });
            }
        }
    }

    Ok(smd)
}

fn process_fbx_node(
    fbx: &SimpleFbx,
    fbx_node: &ObjectTreeNode, mut smd: &mut Smd,
    matrix: &Matrix4<f32>,
    current_bone: Option<&SmdBone>,
) -> Result<(), Error> {
    // Perform node type specific information
    match fbx_node.object.class {
        ObjectType::Geometry(ref geometry) =>
            process_geometry(smd, geometry, matrix, current_bone.unwrap()),
        ObjectType::Model(ref _model) =>
            process_model(fbx, fbx_node, smd, matrix, current_bone)?,
        _ => {
            // Just go straight to the children
            for node in &fbx_node.nodes {
                process_fbx_node(fbx, node, smd, matrix, current_bone)?;
            }
        }
    }

    Ok(())
}

fn process_geometry(smd: &mut Smd, geometry: &Geometry, matrix: &Matrix4<f32>, current_bone: &SmdBone) {
    // Add triangles to parent node
    let tris = geometry.triangles();
    for tri in tris {
        // Turn the vertices in this triangle to SMD vertices
        let mut smd_verts: [SmdVertex; 3] = Default::default();
        for (i, vert) in tri.iter().enumerate() {
            // Multiply the vectors that need to be multiplied
            let pos = matrix * Vector4::new(vert.0[0], vert.0[1], vert.0[2], 1.0);
            let norm = matrix * Vector4::new(vert.1[0], vert.1[1], vert.1[2], 0.0);

            smd_verts[i] = SmdVertex {
                parent_bone: current_bone.id, // This is overwritten by links
                position: pos.truncate().into(),
                normal: norm.truncate().into(),
                uv: vert.2,
                links: vec!(
                    /*Not needed, we aren't using weights anyways so this done by parent_bone
                    SmdLink {
                        bone: bone_id,
                        weight: 1.0,
                    }*/
                )
            };
        }

        // Add the actual SMD triangle
        smd.triangles.push(SmdTriangle {
            material: "layl_test_texture".into(),
            vertices: smd_verts,
        });
    }
}

fn process_model(
    fbx: &SimpleFbx,
    fbx_node: &ObjectTreeNode, smd: &mut Smd,
    matrix: &Matrix4<f32>,
    current_bone: Option<&SmdBone>
) -> Result<(), Error> {
    task_log(format!("Adding model \"{}\" to SMD data", friendly_name(&fbx_node.object.name)));
    let properties = ModelProperties::from_generic(&fbx_node.object.properties);

    // Create a new transformation matrix
    let local_matrix = local_matrices(&properties);

    // Create a new bone
    let new_bone = smd.new_bone(
            &id_name(&fbx_node.object.name).unwrap(),
            current_bone.map(|b| b.id)
        )
        .ok_or_else(|| Error::Task(format!(
            "Bone \"{}\" exists multiple times in the FBX",
            &fbx_node.object.name
        )))?
        .clone(); // Clone needed to avoid a borrow since we need to mut borrow SMD later

    // Set the transformations on this bone
    let (translation, rotation) = calculate_animation_transforms_for(&fbx_node.object);
    let first_frame = SmdAnimationFrameBone {
        // This needs to be derived from the matrix to get the right location
        translation: translation.into(),
        // This can just be directly copied over
        rotation: rotation.into(),
    };
    smd.set_animation(0, new_bone.id, first_frame);

    // Make new matrices for children
    let matrix = matrix * local_matrix;

    // Make sure the child nodes will receive this new bone
    for node in &fbx_node.nodes {
        process_fbx_node(fbx, node, smd, &matrix, Some(&new_bone))?;
    }

    Ok(())
}

/// Returns (Translation, Rotation)
fn calculate_animation_transforms_for(obj: &Object) -> (Vector3<f32>, Vector3<f32>) {
    let properties = ModelProperties::from_generic(&obj.properties);

    // Now get the bone's translation
    let prop_translation: Vector3<_> = properties.translation.into();
    //let prop_rot_offset: Vector3<_> = properties.rotation_offset.into();
    //let prop_rot_pivot: Vector3<_> = properties.rotation_pivot.into();
    let translation = prop_translation;

    // We want the rotation, but we've got multiple rotations, so combine them
    let pre_rotation = Quaternion::from(Euler::new(
        Deg(properties.pre_rotation[0]), Deg(properties.pre_rotation[1]), Deg(properties.pre_rotation[2])
    ));
    let rotation = Quaternion::from(Euler::new(
        Deg(properties.rotation[0]), Deg(properties.rotation[1]), Deg(properties.rotation[2])
    ));
    let post_rotation = Quaternion::from(Euler::new(
        Deg(properties.post_rotation[0]), Deg(properties.post_rotation[1]), Deg(properties.post_rotation[2])
    ));

    let total_rotation = Euler::from(post_rotation * rotation * pre_rotation);
    let rotation = Vector3::new(
        total_rotation.x.0,
        total_rotation.y.0,
        total_rotation.z.0,
    );

    (translation, rotation)
}

fn local_matrices(properties: &ModelProperties) -> Matrix4<f32> {
    // Create various matrices
    let rotation_offset = properties.rotation_offset.into();
    let rotation_offset_mat = Matrix4::from_translation(rotation_offset);
    let rotation_pivot: Vector3<_> = properties.rotation_pivot.into();
    let rotation_pivot_mat = Matrix4::from_translation(rotation_pivot);

    let pre_rotation = euler_rotation_to_matrix(properties.pre_rotation);
    let rotation = euler_rotation_to_matrix(properties.rotation);
    let post_rotation = euler_rotation_to_matrix(properties.post_rotation);

    let scale_offset = properties.scale_offset.into();
    let scale_offset_mat = Matrix4::from_translation(scale_offset);
    let scale_pivot: Vector3<_> = properties.scale_pivot.into();
    let scale_pivot_mat = Matrix4::from_translation(scale_pivot);

    let scale = Matrix4::from_nonuniform_scale(
        properties.scale[0],
        properties.scale[1],
        properties.scale[2]
    );

    let local_matrix_for_vertices =
        Matrix4::from_translation(properties.translation.into()) *

        // Rotation
        rotation_offset_mat *
        rotation_pivot_mat *
        pre_rotation *
        rotation *
        post_rotation *
        rotation_pivot_mat.invert().unwrap() *

        // Scale
        scale_offset_mat *
        scale_pivot_mat *
        scale *
        scale_pivot_mat.invert().unwrap();

    local_matrix_for_vertices
}

fn euler_rotation_to_matrix(rot_degs: [f32; 3]) -> Matrix4<f32> {
    Matrix4::from_angle_z(Deg(rot_degs[2])) *
    Matrix4::from_angle_y(Deg(rot_degs[1])) *
    Matrix4::from_angle_x(Deg(rot_degs[0]))
}
<commit_msg>Once again fix pivot positions this time the Smart Way<commit_after>use std::fs::File;
use std::io::{BufReader};
use std::path::PathBuf;

use cgmath::{Matrix4, Deg, Vector4, SquareMatrix, Vector3, Euler, Quaternion};
use soto::task::{task_log};
use soto::Error;
use sotolib_fbx::{RawFbx, id_name, friendly_name, ObjectTreeNode};
use sotolib_fbx::animation::{Animation};
use sotolib_fbx::simple::{Object, SimpleFbx, ObjectType, ModelProperties, Geometry};
use sotolib_smd::{Smd, SmdVertex, SmdTriangle, SmdAnimationFrameBone, SmdBone};

pub fn create_reference_smd(fbx: &PathBuf) -> Result<Smd, Error> {
    // Read in the fbx we got told to convert
    let file = BufReader::new(File::open(&fbx).unwrap());
    let fbx = SimpleFbx::from_raw(&RawFbx::parse(file).unwrap()).unwrap();
    let fbx_tree = ObjectTreeNode::from_simple(&fbx);

    // Go over all FBX root nodes and turn them into SMD data
    let mut smd = Smd::new();
    process_fbx_node(
        &fbx,
        &fbx_tree, &mut smd,
        &Matrix4::identity(),
        None
    )?;

    Ok(smd)
}

pub fn create_animation_smd(ref_smd: &Smd, fbx: &PathBuf) -> Result<Smd, Error> {
    // Read in the fbx we got told to convert
    let file = BufReader::new(File::open(&fbx).unwrap());
    let mut fbx = SimpleFbx::from_raw(&RawFbx::parse(file).unwrap()).unwrap();

    // Read in the animation data itself
    let animation = Animation::from_simple(&fbx).unwrap();

    // Count and log frames
    let frame_count = animation.frame_count(&fbx);
    task_log(format!("Animation has {} frames", frame_count));

    // Copy over every bone to the new animation SMD
    let mut smd = Smd::new();
    for bone in &ref_smd.bones {
        smd.bones.push(bone.clone());
    }

    // Finally, turn the animation data into bone positions in the SMD
    for frame in 0..frame_count {
        // First transform the FBX for this frame
        animation.transform_fbx_to_frame(&mut fbx, frame);

        // Now go over all models
        for (_, model) in fbx.objects.iter().filter(|&(_, o)| o.class.type_name() == "Model") {
            // For this model, look up the matching BoneId in the reference SMD
            if let Some(bone_id) = ref_smd.id_of_bone(&id_name(&model.name).unwrap()) {
                // Now that we have a model and a bone, we need the current translation and rotation
                // for the model
                let (translation, rotation) = calculate_animation_transforms_for(&fbx, model);

                // And now that we have those, finally add the bone data to the animation SMD
                smd.set_animation(frame, bone_id, SmdAnimationFrameBone {
                    translation: translation.into(),
                    rotation: rotation.into(),
                });
            }
        }
    }

    Ok(smd)
}

fn process_fbx_node(
    fbx: &SimpleFbx,
    fbx_node: &ObjectTreeNode, mut smd: &mut Smd,
    matrix: &Matrix4<f32>,
    current_bone: Option<&SmdBone>,
) -> Result<(), Error> {
    // Perform node type specific information
    match fbx_node.object.class {
        ObjectType::Geometry(ref geometry) =>
            process_geometry(smd, geometry, matrix, current_bone.unwrap()),
        ObjectType::Model(ref _model) =>
            process_model(fbx, fbx_node, smd, matrix, current_bone)?,
        _ => {
            // Just go straight to the children
            for node in &fbx_node.nodes {
                process_fbx_node(fbx, node, smd, matrix, current_bone)?;
            }
        }
    }

    Ok(())
}

fn process_geometry(smd: &mut Smd, geometry: &Geometry, matrix: &Matrix4<f32>, current_bone: &SmdBone) {
    // Add triangles to parent node
    let tris = geometry.triangles();
    for tri in tris {
        // Turn the vertices in this triangle to SMD vertices
        let mut smd_verts: [SmdVertex; 3] = Default::default();
        for (i, vert) in tri.iter().enumerate() {
            // Multiply the vectors that need to be multiplied
            let pos = matrix * Vector4::new(vert.0[0], vert.0[1], vert.0[2], 1.0);
            let norm = matrix * Vector4::new(vert.1[0], vert.1[1], vert.1[2], 0.0);

            smd_verts[i] = SmdVertex {
                parent_bone: current_bone.id, // This is overwritten by links
                position: pos.truncate().into(),
                normal: norm.truncate().into(),
                uv: vert.2,
                links: vec!(
                    /*Not needed, we aren't using weights anyways so this done by parent_bone
                    SmdLink {
                        bone: bone_id,
                        weight: 1.0,
                    }*/
                )
            };
        }

        // Add the actual SMD triangle
        smd.triangles.push(SmdTriangle {
            material: "layl_test_texture".into(),
            vertices: smd_verts,
        });
    }
}

fn process_model(
    fbx: &SimpleFbx,
    fbx_node: &ObjectTreeNode, smd: &mut Smd,
    matrix: &Matrix4<f32>,
    current_bone: Option<&SmdBone>
) -> Result<(), Error> {
    task_log(format!("Adding model \"{}\" to SMD data", friendly_name(&fbx_node.object.name)));
    let properties = ModelProperties::from_generic(&fbx_node.object.properties);

    // Create a new transformation matrix
    let local_matrix = local_matrices(&properties);

    // Create a new bone
    let new_bone = smd.new_bone(
            &id_name(&fbx_node.object.name).unwrap(),
            current_bone.map(|b| b.id)
        )
        .ok_or_else(|| Error::Task(format!(
            "Bone \"{}\" exists multiple times in the FBX",
            &fbx_node.object.name
        )))?
        .clone(); // Clone needed to avoid a borrow since we need to mut borrow SMD later

    // Set the transformations on this bone
    let (translation, rotation) = calculate_animation_transforms_for(&fbx, &fbx_node.object);
    let first_frame = SmdAnimationFrameBone {
        // This needs to be derived from the matrix to get the right location
        translation: translation.into(),
        // This can just be directly copied over
        rotation: rotation.into(),
    };
    smd.set_animation(0, new_bone.id, first_frame);

    // Make new matrices for children
    let matrix = matrix * local_matrix;

    // Make sure the child nodes will receive this new bone
    for node in &fbx_node.nodes {
        process_fbx_node(fbx, node, smd, &matrix, Some(&new_bone))?;
    }

    Ok(())
}

/// Returns (Translation, Rotation)
fn calculate_animation_transforms_for(fbx: &SimpleFbx, obj: &Object) -> (Vector3<f32>, Vector3<f32>) {
    let properties = ModelProperties::from_generic(&obj.properties);

    // Get the bone's translation
    let parent_after_rot_translation = calculate_parent_after_rot_translation(fbx, obj);
    let prop_translation: Vector3<_> = properties.translation.into();
    let prop_rot_offset: Vector3<_> = properties.rotation_offset.into();
    let prop_rot_pivot: Vector3<_> = properties.rotation_pivot.into();
    let translation = parent_after_rot_translation + prop_translation + prop_rot_offset + prop_rot_pivot;

    // We want the rotation, but we've got multiple rotations, so combine them
    let pre_rotation = Quaternion::from(Euler::new(
        Deg(properties.pre_rotation[0]), Deg(properties.pre_rotation[1]), Deg(properties.pre_rotation[2])
    ));
    let rotation = Quaternion::from(Euler::new(
        Deg(properties.rotation[0]), Deg(properties.rotation[1]), Deg(properties.rotation[2])
    ));
    let post_rotation = Quaternion::from(Euler::new(
        Deg(properties.post_rotation[0]), Deg(properties.post_rotation[1]), Deg(properties.post_rotation[2])
    ));

    let total_rotation = Euler::from(post_rotation * rotation * pre_rotation);
    let rotation = Vector3::new(
        total_rotation.x.0,
        total_rotation.y.0,
        total_rotation.z.0,
    );

    (translation, rotation)
}

fn calculate_parent_after_rot_translation(fbx: &SimpleFbx, obj: &Object) -> Vector3<f32> {
    // First actually get the parent's model data
    let parent_obj = if let Some(v) = fbx.parent_of(obj.id) {
        if v == 0 {
            // At root, no extra translation
            return Vector3::new(0.0, 0.0, 0.0)
        }
        v
    } else {
        // No parent, no extra translation
        return Vector3::new(0.0, 0.0, 0.0)
    };
    let props = ModelProperties::from_generic(&fbx.objects[&parent_obj].properties);

    // Now add up all the translations applied after rotation
    let rotation_pivot: Vector3<_> = props.rotation_pivot.into();
    let scale_offset: Vector3<_> = props.scale_offset.into();
    let translation = -rotation_pivot + scale_offset;

    translation
}

fn local_matrices(properties: &ModelProperties) -> Matrix4<f32> {
    // Create various matrices
    let rotation_offset = properties.rotation_offset.into();
    let rotation_offset_mat = Matrix4::from_translation(rotation_offset);
    let rotation_pivot: Vector3<_> = properties.rotation_pivot.into();
    let rotation_pivot_mat = Matrix4::from_translation(rotation_pivot);

    let pre_rotation = euler_rotation_to_matrix(properties.pre_rotation);
    let rotation = euler_rotation_to_matrix(properties.rotation);
    let post_rotation = euler_rotation_to_matrix(properties.post_rotation);

    let scale_offset = properties.scale_offset.into();
    let scale_offset_mat = Matrix4::from_translation(scale_offset);
    let scale_pivot: Vector3<_> = properties.scale_pivot.into();
    let scale_pivot_mat = Matrix4::from_translation(scale_pivot);

    let scale = Matrix4::from_nonuniform_scale(
        properties.scale[0],
        properties.scale[1],
        properties.scale[2]
    );

    let local_matrix_for_vertices =
        Matrix4::from_translation(properties.translation.into()) *

        // Rotation
        rotation_offset_mat *
        rotation_pivot_mat *
        pre_rotation *
        rotation *
        post_rotation *
        rotation_pivot_mat.invert().unwrap() *

        // Scale
        scale_offset_mat *
        scale_pivot_mat *
        scale *
        scale_pivot_mat.invert().unwrap();

    local_matrix_for_vertices
}

fn euler_rotation_to_matrix(rot_degs: [f32; 3]) -> Matrix4<f32> {
    Matrix4::from_angle_z(Deg(rot_degs[2])) *
    Matrix4::from_angle_y(Deg(rot_degs[1])) *
    Matrix4::from_angle_x(Deg(rot_degs[0]))
}
<|endoftext|>