<commit_before>extern crate anyhow;
extern crate arrayvec;
extern crate dpdk;
extern crate log;
extern crate simple_logger;

use anyhow::Result;
use arrayvec::*;
use dpdk::eal::*;
use log::{debug, info};
use std::env;

fn sender(eal: Eal, mpool: MPool, tx_queue: TxQ) {
    let tx_port = tx_queue.port();
    info!("Start TX from {:?}", tx_port.mac_addr());

    // Wait for the link to be connected.
    while !tx_port.is_link_up() {
        eal.pause();
    }
    info!("Link is up {:?}", tx_port.mac_addr());

    let mut pkts = ArrayVec::<[Packet; DEFAULT_TX_BURST]>::new();
    // Safety: packet is created and transmitted before `mpool` is destroyed.
    unsafe { mpool.alloc_bulk(&mut pkts) };

    for pkt in &mut pkts {
        // Prepare toy arp request packets
        let pkt_buf = pkt.buffer_mut();

        // Fill Ethernet
        // TODO Use https://doc.rust-lang.org/beta/std/primitive.slice.html#method.fill
        pkt_buf[0..6].copy_from_slice(&[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF]); // Dst MAC
        pkt_buf[6..12].copy_from_slice(&tx_port.mac_addr()); // Src MAC
        pkt_buf[12..14].copy_from_slice(&[0x08, 0x06]); // Ethertype: ARP

        // Fill ARP
        pkt_buf[14..16].copy_from_slice(&[0x00, 0x01]); // HTYPE: Ethernet
        pkt_buf[16..18].copy_from_slice(&[0x08, 0x00]); // PTYPE: IP
        pkt_buf[18..20].copy_from_slice(&[0x06, 0x04]); // HLEN: 6, PLEN: 4
        pkt_buf[20..22].copy_from_slice(&[0x00, 0x01]); // OPER: request (1)
        pkt_buf[22..28].copy_from_slice(&tx_port.mac_addr()); // SHA (6byte)
        pkt_buf[28..32].copy_from_slice(&[0x10, 0x00, 0x00, 0x02]); // SPA (10.0.0.2)
        pkt_buf[32..38].copy_from_slice(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00]); // THA (6byte)
        pkt_buf[38..42].copy_from_slice(&[0x10, 0x00, 0x00, 0x03]); // THA (10.0.0.3)

        pkt.set_len(42);
    }
    // Send packet
    tx_queue.tx(&mut pkts);

    // Wait for pkts to be transmitted
    while tx_port.get_stat().opackets as usize > DEFAULT_TX_BURST {
        eal.pause();
    }

    info!("TX finished. {:?}", tx_port.get_stat());

    // Safety: mpool must not be deallocated before TxQ is destroyed.
}

fn receiver(eal: Eal, rx_queue: RxQ) {
    let rx_port = rx_queue.port();
    info!("RX started at {:?}", rx_port.mac_addr());

    // We will try to collect every TX packets.
    // We will collect all sent packets and additional background packets.
    // Thus we need 2 * TX_BURST to collect everything.
    let mut pkts = ArrayVec::<[Packet; DEFAULT_TX_BURST * 2]>::new();
    loop {
        rx_queue.rx(&mut pkts);
        if pkts.len() >= DEFAULT_TX_BURST {
            break;
        }
        eal.pause();
    }
    info!("RX finished. {:?}", rx_port.get_stat());
}

/// Note: this test function only works with `sudo target/debug/dpdk_test -c 1`
fn main() -> Result<()> {
    simple_logger::init().unwrap();
    let mut args: Vec<String> = env::args().collect();
    let eal = Eal::new(&mut args).unwrap();
    debug!("TSC Hz: {}", eal.get_tsc_hz());
    debug!("TSC Cycles: {}", eal.get_tsc_cycles());
    debug!("Random: {}", eal.rand());

    // Create a `MPool` to create packets.
    let default_mpool = eal.create_mpool(
        "default_tx_pool",
        DEFAULT_RX_POOL_SIZE,
        DEFAULT_RX_PER_CORE_CACHE,
        DEFAULT_PACKET_DATA_LENGTH,
        None,
    );

    let threads = eal
        .setup(Affinity::Full, Affinity::Full)?
        .into_iter()
        .map(|(lcore, rxs, txs)| {
            let local_eal = eal.clone();
            let local_mpool = default_mpool.clone();
            lcore.launch(move || {
                match lcore.into() {
                    // Core 0 action: TX packets to txq[0]
                    0 => {
                        sender(local_eal.clone(), local_mpool, txs[0].clone());
                        receiver(local_eal, rxs[1].clone());
                        true
                    }
                    // Otherwise, do nothing
                    _ => true,
                }
            })
        })
        .collect::<Vec<_>>();
    let ret = threads.into_iter().map(|x| x.join().unwrap()).all(|x| x);
    assert_eq!(ret, true);
    Ok(())
}
<commit_msg>쌍방에 모두 추가<commit_after>extern crate anyhow;
extern crate arrayvec;
extern crate dpdk;
extern crate log;
extern crate simple_logger;

use anyhow::Result;
use arrayvec::*;
use dpdk::eal::*;
use log::{debug, info};
use std::env;

fn sender(eal: Eal, mpool: MPool, tx_queue: TxQ) {
    let tx_port = tx_queue.port();
    info!("Start TX from {:?}", tx_port.mac_addr());

    // Wait for the link to be connected.
    while !tx_port.is_link_up() {
        eal.pause();
    }
    info!("TX Link is up {:?}", tx_port.mac_addr());

    let mut pkts = ArrayVec::<[Packet; DEFAULT_TX_BURST]>::new();
    // Safety: packet is created and transmitted before `mpool` is destroyed.
    unsafe { mpool.alloc_bulk(&mut pkts) };

    for pkt in &mut pkts {
        // Prepare toy arp request packets
        let pkt_buf = pkt.buffer_mut();

        // Fill Ethernet
        // TODO Use https://doc.rust-lang.org/beta/std/primitive.slice.html#method.fill
        pkt_buf[0..6].copy_from_slice(&[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF]); // Dst MAC
        pkt_buf[6..12].copy_from_slice(&tx_port.mac_addr()); // Src MAC
        pkt_buf[12..14].copy_from_slice(&[0x08, 0x06]); // Ethertype: ARP

        // Fill ARP
        pkt_buf[14..16].copy_from_slice(&[0x00, 0x01]); // HTYPE: Ethernet
        pkt_buf[16..18].copy_from_slice(&[0x08, 0x00]); // PTYPE: IP
        pkt_buf[18..20].copy_from_slice(&[0x06, 0x04]); // HLEN: 6, PLEN: 4
        pkt_buf[20..22].copy_from_slice(&[0x00, 0x01]); // OPER: request (1)
        pkt_buf[22..28].copy_from_slice(&tx_port.mac_addr()); // SHA (6byte)
        pkt_buf[28..32].copy_from_slice(&[0x10, 0x00, 0x00, 0x02]); // SPA (10.0.0.2)
        pkt_buf[32..38].copy_from_slice(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00]); // THA (6byte)
        pkt_buf[38..42].copy_from_slice(&[0x10, 0x00, 0x00, 0x03]); // THA (10.0.0.3)

        pkt.set_len(42);
    }
    // Send packet
    tx_queue.tx(&mut pkts);

    // Wait for pkts to be transmitted
    while tx_port.get_stat().opackets as usize > DEFAULT_TX_BURST {
        eal.pause();
    }

    info!("TX finished. {:?}", tx_port.get_stat());

    // Safety: mpool must not be deallocated before TxQ is destroyed.
}

fn receiver(eal: Eal, rx_queue: RxQ) {
    let rx_port = rx_queue.port();
    info!("RX started at {:?}", rx_port.mac_addr());

    // Wait for the link to be connected.
    while !rx_port.is_link_up() {
        eal.pause();
    }
    info!("RX Link is up {:?}", rx_port.mac_addr());

    // We will try to collect every TX packets.
    // We will collect all sent packets and additional background packets.
    // Thus we need 2 * TX_BURST to collect everything.
    let mut pkts = ArrayVec::<[Packet; DEFAULT_TX_BURST * 2]>::new();
    loop {
        rx_queue.rx(&mut pkts);
        if pkts.len() >= DEFAULT_TX_BURST {
            break;
        }
        eal.pause();
    }
    info!("RX finished. {:?}", rx_port.get_stat());
}

/// Note: this test function only works with `sudo target/debug/dpdk_test -c 1`
fn main() -> Result<()> {
    simple_logger::init().unwrap();
    let mut args: Vec<String> = env::args().collect();
    let eal = Eal::new(&mut args).unwrap();
    debug!("TSC Hz: {}", eal.get_tsc_hz());
    debug!("TSC Cycles: {}", eal.get_tsc_cycles());
    debug!("Random: {}", eal.rand());

    // Create a `MPool` to create packets.
    let default_mpool = eal.create_mpool(
        "default_tx_pool",
        DEFAULT_RX_POOL_SIZE,
        DEFAULT_RX_PER_CORE_CACHE,
        DEFAULT_PACKET_DATA_LENGTH,
        None,
    );

    let threads = eal
        .setup(Affinity::Full, Affinity::Full)?
        .into_iter()
        .map(|(lcore, rxs, txs)| {
            let local_eal = eal.clone();
            let local_mpool = default_mpool.clone();
            lcore.launch(move || {
                match lcore.into() {
                    // Core 0 action: TX packets to txq[0]
                    0 => {
                        sender(local_eal.clone(), local_mpool, txs[0].clone());
                        receiver(local_eal, rxs[1].clone());
                        true
                    }
                    // Otherwise, do nothing
                    _ => true,
                }
            })
        })
        .collect::<Vec<_>>();
    let ret = threads.into_iter().map(|x| x.join().unwrap()).all(|x| x);
    assert_eq!(ret, true);
    Ok(())
}
<|endoftext|>


<commit_before>
extern crate nccl;

use nccl::value::Value;
use nccl::pair::Pair;

#[test]
fn from_vec() {
    assert_eq!(Value::from(vec![1, 2, 3]),
               Value::List(vec![Value::Integer(1),
                                Value::Integer(2),
                                Value::Integer(3)]));
}

#[test]
fn add() {
    let mut p = Pair::new("key");
    p.add("value");
    assert_eq!(p, Pair {
        key: "key".into(),
        value: vec![Pair {
            key: "value".into(),
            value: vec![]
        }]
    });
}

#[test]
fn get() {
    let mut p = Pair::new("key");
    p.add("value");
    p.add("testaroni");
    p.get_mut("value".into()).unwrap().add("Hello!");
    let v = p.get("value".into()).unwrap();
    assert_eq!(v, &Pair {
        key: "value".into(),
        value: vec![Pair {
            key: "Hello!".into(),
            value: vec![]
        }]
    });
}

#[test]
fn multi() {
    let mut p = Pair::new("key");
    p.add("value");
    p.add("dj khaled");
    p.get_mut("value".into()).unwrap().add("yes!");
    p.get_mut("dj khaled".into()).unwrap().add("you smart");
    p.get_mut("dj khaled".into()).unwrap().add("you loyal");
    assert_eq!(p, Pair {
        key: "key".into(),
        value: vec![Pair {
            key: "value".into(),
            value: vec![Pair {
                key: "yes!".into(),
                value: vec![]
            }]
        }, Pair {
            key: "dj khaled".into(),
            value: vec![Pair {
                key: "you smart".into(),
                value: vec![]
            }, Pair{
                key: "you loyal".into(),
                value: vec![]
            }]
        }]
    });
}

#[test]
fn index() {
    let mut p = Pair::new("key");
    p.add("hello");
    p["hello".into()].add("world");
    p["hello".into()]["world".into()].add("what's");
    p["hello".into()]["world".into()]["what's".into()].add("up?");
    assert_eq!(p, Pair {
        key: "key".into(),
        value: vec![Pair {
            key: "hello".into(),
            value: vec![Pair {
                key: "world".into(),
                value: vec![Pair {
                    key: "what's".into(),
                    value: vec![Pair {
                        key: "up?".into(),
                        value: vec![]
                    }]
                }]
            }]
        }]
    });
}

<commit_msg>update tests<commit_after>
extern crate nccl;

use nccl::value::Value;
use nccl::pair::Pair;

#[test]
fn from_vec() {
    assert_eq!(Value::from(vec![1, 2, 3]),
               Value::List(vec![Value::Integer(1),
                                Value::Integer(2),
                                Value::Integer(3)]));
}

#[test]
fn add() {
    let mut p = Pair::new("key");
    p.add("value");
    assert_eq!(p, Pair {
        key: "key".into(),
        value: vec![Pair {
            key: "value".into(),
            value: vec![]
        }]
    });
}

#[test]
fn get() {
    let mut p = Pair::new("key");
    p.add("value");
    p.add("testaroni");
    p.get_mut("value".into()).unwrap().add("Hello!");
    let v = p.get("value".into()).unwrap();
    assert_eq!(v, &Pair {
        key: "value".into(),
        value: vec![Pair {
            key: "Hello!".into(),
            value: vec![]
        }]
    });
}

#[test]
fn multi() {
    let mut p = Pair::new("key");
    p.add("value");
    p.add("dj khaled");
    p.get_mut("value".into()).unwrap().add("yes!");
    p.get_mut("dj khaled".into()).unwrap().add("you smart");
    p.get_mut("dj khaled".into()).unwrap().add("you loyal");
    assert_eq!(p, Pair {
        key: "key".into(),
        value: vec![Pair {
            key: "value".into(),
            value: vec![Pair {
                key: "yes!".into(),
                value: vec![]
            }]
        }, Pair {
            key: "dj khaled".into(),
            value: vec![Pair {
                key: "you smart".into(),
                value: vec![]
            }, Pair{
                key: "you loyal".into(),
                value: vec![]
            }]
        }]
    });
}

#[test]
fn index() {
    let mut p = Pair::new("key");
    p.add("hello");
    p["hello"].add("world");
    p["hello"]["world"].add("what's");
    p["hello"]["world"]["what's"].add("up?");
    assert_eq!(p, Pair {
        key: "key".into(),
        value: vec![Pair {
            key: "hello".into(),
            value: vec![Pair {
                key: "world".into(),
                value: vec![Pair {
                    key: "what's".into(),
                    value: vec![Pair {
                        key: "up?".into(),
                        value: vec![]
                    }]
                }]
            }]
        }]
    });
}

<|endoftext|>


<commit_before>#[feature(globs)];
#[feature(macro_rules)];
#[feature(managed_boxes)];

#[link_args="-lwxc"];

extern mod native;
extern mod wx;

use std::libc::c_void;

use wx::_unsafe::*;
use wx::defs::*;
use wx::base::*;
use wx::core::*;

mod macros;

wxApp!(wx_main)

extern "C"
fn wx_main() {
    let frame = MyFrame::new();
    frame.show();
    frame.raise();
}

struct MyFrame(@wxFrame);
impl MyFrame {
    fn new() -> MyFrame {
        let frame = wxFrame::new(wxWindow::null(), wxID_ANY, "Hello, wxRust!", -1, -1, -1, -1, wxDEFAULT_FRAME_STYLE);
        let menubar = MyMenuBar::new();
        frame.setMenuBar(menubar.asMenuBar());

        MyButton::new(frame);

        MyFrame(frame)
    }
}

struct MyMenuBar { handle: @wxMenuBar }
impl MyMenuBar {
    fn new() -> MyMenuBar {
        let menubar = wxMenuBar::new(0);

        let fileMenu = wxMenu::new("", 0);
        let fileNew = wxMenuItem::newEx(wxID_ANY, "New", "Create a new file.", 0, wxMenu::null());
        fileMenu.appendItem(fileNew);

        menubar.append(fileMenu, "File");
        MyMenuBar { handle: menubar }
    }
    fn asMenuBar(&self) -> @wxMenuBar {
        return self.handle;
    }
}

extern "C"
fn MyButton_clicked(fun: *mut c_void, data: *mut c_void, evt: *mut c_void) {
    if (evt == 0 as *mut c_void) {
        // Comes here when the target widget is destroyed.
        return;
    }

    println("hello!");
    let parent = wxWindow::from(data);
    let msgDlg = wxMessageDialog::new(parent, "Pushed!!", "The Button", wxOK);
    msgDlg.showModal();
}

struct MyButton(@wxButton);
impl MyButton {
    fn new<T: _wxWindow>(parent: &T) -> MyButton {
        let button = wxButton::new(parent, wxID_ANY, "Push me!", 10, 10, 50, 30, 0);
        let closure = wxClosure::new(MyButton_clicked as *mut c_void, parent.handle());
        unsafe {
            button.connect(wxID_ANY, wxID_ANY, expEVT_COMMAND_BUTTON_CLICKED(), closure.handle());
        }

        MyButton(button)
    }
}
<commit_msg>Don't rely on auto-dereferencing to fix build.<commit_after>#[feature(globs)];
#[feature(macro_rules)];
#[feature(managed_boxes)];

#[link_args="-lwxc"];

extern mod native;
extern mod wx;

use std::libc::c_void;

use wx::_unsafe::*;
use wx::defs::*;
use wx::base::*;
use wx::core::*;

mod macros;

wxApp!(wx_main)

extern "C"
fn wx_main() {
    let frame = MyFrame::new();
    frame.asFrame().show();
    frame.asFrame().raise();
}

struct MyFrame { base: @wxFrame }
impl MyFrame {
    fn new() -> MyFrame {
        let frame = wxFrame::new(wxWindow::null(), wxID_ANY, "Hello, wxRust!", -1, -1, -1, -1, wxDEFAULT_FRAME_STYLE);
        let menubar = MyMenuBar::new();
        frame.setMenuBar(menubar.asMenuBar());

        MyButton::new(frame);

        MyFrame { base: frame }
    }
    fn asFrame(&self) -> @wxFrame {
        return self.base;
    }
}

struct MyMenuBar { base: @wxMenuBar }
impl MyMenuBar {
    fn new() -> MyMenuBar {
        let menubar = wxMenuBar::new(0);

        let fileMenu = wxMenu::new("", 0);
        let fileNew = wxMenuItem::newEx(wxID_ANY, "New", "Create a new file.", 0, wxMenu::null());
        fileMenu.appendItem(fileNew);

        menubar.append(fileMenu, "File");
        MyMenuBar { base: menubar }
    }
    fn asMenuBar(&self) -> @wxMenuBar {
        return self.base;
    }
}

extern "C"
fn MyButton_clicked(fun: *mut c_void, data: *mut c_void, evt: *mut c_void) {
    if (evt == 0 as *mut c_void) {
        // Comes here when the target widget is destroyed.
        return;
    }

    println("hello!");
    let parent = wxWindow::from(data);
    let msgDlg = wxMessageDialog::new(parent, "Pushed!!", "The Button", wxOK);
    msgDlg.showModal();
}

struct MyButton(@wxButton);
impl MyButton {
    fn new<T: _wxWindow>(parent: &T) -> MyButton {
        let button = wxButton::new(parent, wxID_ANY, "Push me!", 10, 10, 50, 30, 0);
        let closure = wxClosure::new(MyButton_clicked as *mut c_void, parent.handle());
        unsafe {
            button.connect(wxID_ANY, wxID_ANY, expEVT_COMMAND_BUTTON_CLICKED(), closure.handle());
        }

        MyButton(button)
    }
}
<|endoftext|>


<commit_before>use super::{Open, Sink, SinkAsBytes, SinkError, SinkResult};
use crate::config::AudioFormat;
use crate::convert::Converter;
use crate::decoder::AudioPacket;
use crate::{NUM_CHANNELS, SAMPLE_RATE};
use alsa::device_name::HintIter;
use alsa::pcm::{Access, Format, HwParams, PCM};
use alsa::{Direction, ValueOr};
use std::cmp::min;
use std::process::exit;
use std::time::Duration;
use thiserror::Error;

// 0.5 sec buffer.
const PERIOD_TIME: Duration = Duration::from_millis(100);
const BUFFER_TIME: Duration = Duration::from_millis(500);

#[derive(Debug, Error)]
enum AlsaError {
    #[error("<AlsaSink> Device {device} Unsupported Format {alsa_format:?} ({format:?}), {e}")]
    UnsupportedFormat {
        device: String,
        alsa_format: Format,
        format: AudioFormat,
        e: alsa::Error,
    },

    #[error("<AlsaSink> Device {device} Unsupported Channel Count {channel_count}, {e}")]
    UnsupportedChannelCount {
        device: String,
        channel_count: u8,
        e: alsa::Error,
    },

    #[error("<AlsaSink> Device {device} Unsupported Sample Rate {samplerate}, {e}")]
    UnsupportedSampleRate {
        device: String,
        samplerate: u32,
        e: alsa::Error,
    },

    #[error("<AlsaSink> Device {device} Unsupported Access Type RWInterleaved, {e}")]
    UnsupportedAccessType { device: String, e: alsa::Error },

    #[error("<AlsaSink> Device {device} May be Invalid, Busy, or Already in Use, {e}")]
    PcmSetUp { device: String, e: alsa::Error },

    #[error("<AlsaSink> Failed to Drain PCM Buffer, {0}")]
    DrainFailure(alsa::Error),

    #[error("<AlsaSink> {0}")]
    OnWrite(alsa::Error),

    #[error("<AlsaSink> Hardware, {0}")]
    HwParams(alsa::Error),

    #[error("<AlsaSink> Software, {0}")]
    SwParams(alsa::Error),

    #[error("<AlsaSink> PCM, {0}")]
    Pcm(alsa::Error),

    #[error("<AlsaSink> Could Not Parse Ouput Name(s) and/or Description(s)")]
    Parsing,

    #[error("<AlsaSink>")]
    NotConnected,
}

impl From<AlsaError> for SinkError {
    fn from(e: AlsaError) -> SinkError {
        use AlsaError::*;
        let es = e.to_string();
        match e {
            DrainFailure(_) | OnWrite(_) => SinkError::OnWrite(es),
            PcmSetUp { .. } => SinkError::ConnectionRefused(es),
            NotConnected => SinkError::NotConnected(es),
            _ => SinkError::InvalidParams(es),
        }
    }
}

pub struct AlsaSink {
    pcm: Option<PCM>,
    format: AudioFormat,
    device: String,
    period_buffer: Vec<u8>,
}

fn list_outputs() -> SinkResult<()> {
    println!("Listing available Alsa outputs:");
    for t in &["pcm", "ctl", "hwdep"] {
        println!("{} devices:", t);

        let i = HintIter::new_str(None, &t).map_err(|_| AlsaError::Parsing)?;

        for a in i {
            if let Some(Direction::Playback) = a.direction {
                // mimic aplay -L
                let name = a.name.ok_or(AlsaError::Parsing)?;
                let desc = a.desc.ok_or(AlsaError::Parsing)?;

                println!("{}\n\t{}\n", name, desc.replace("\n", "\n\t"));
            }
        }
    }

    Ok(())
}

fn open_device(dev_name: &str, format: AudioFormat) -> SinkResult<(PCM, usize)> {
    let pcm = PCM::new(dev_name, Direction::Playback, false).map_err(|e| AlsaError::PcmSetUp {
        device: dev_name.to_string(),
        e,
    })?;

    let alsa_format = match format {
        AudioFormat::F64 => Format::float64(),
        AudioFormat::F32 => Format::float(),
        AudioFormat::S32 => Format::s32(),
        AudioFormat::S24 => Format::s24(),
        AudioFormat::S16 => Format::s16(),

        #[cfg(target_endian = "little")]
        AudioFormat::S24_3 => Format::S243LE,
        #[cfg(target_endian = "big")]
        AudioFormat::S24_3 => Format::S243BE,
    };

    let bytes_per_period = {
        let hwp = HwParams::any(&pcm).map_err(AlsaError::HwParams)?;

        hwp.set_access(Access::RWInterleaved)
            .map_err(|e| AlsaError::UnsupportedAccessType {
                device: dev_name.to_string(),
                e,
            })?;

        hwp.set_format(alsa_format)
            .map_err(|e| AlsaError::UnsupportedFormat {
                device: dev_name.to_string(),
                alsa_format,
                format,
                e,
            })?;

        hwp.set_rate(SAMPLE_RATE, ValueOr::Nearest).map_err(|e| {
            AlsaError::UnsupportedSampleRate {
                device: dev_name.to_string(),
                samplerate: SAMPLE_RATE,
                e,
            }
        })?;

        hwp.set_channels(NUM_CHANNELS as u32)
            .map_err(|e| AlsaError::UnsupportedChannelCount {
                device: dev_name.to_string(),
                channel_count: NUM_CHANNELS,
                e,
            })?;

        hwp.set_buffer_time_near(BUFFER_TIME.as_micros() as u32, ValueOr::Nearest)
            .map_err(AlsaError::HwParams)?;

        hwp.set_period_time_near(PERIOD_TIME.as_micros() as u32, ValueOr::Nearest)
            .map_err(AlsaError::HwParams)?;

        pcm.hw_params(&hwp).map_err(AlsaError::Pcm)?;

        let swp = pcm.sw_params_current().map_err(AlsaError::Pcm)?;

        // Don't assume we got what we wanted. Ask to make sure.
        let frames_per_period = hwp.get_period_size().map_err(AlsaError::HwParams)?;

        let frames_per_buffer = hwp.get_buffer_size().map_err(AlsaError::HwParams)?;

        swp.set_start_threshold(frames_per_buffer - frames_per_period)
            .map_err(AlsaError::SwParams)?;

        pcm.sw_params(&swp).map_err(AlsaError::Pcm)?;

        trace!("Frames per Buffer: {:?}", frames_per_buffer);
        trace!("Frames per Period: {:?}", frames_per_period);

        // Let ALSA do the math for us.
        pcm.frames_to_bytes(frames_per_period) as usize
    };

    trace!("Period Buffer size in bytes: {:?}", bytes_per_period);

    Ok((pcm, bytes_per_period))
}

impl Open for AlsaSink {
    fn open(device: Option<String>, format: AudioFormat) -> Self {
        let name = match device.as_deref() {
            Some("?") => match list_outputs() {
                Ok(_) => {
                    exit(0);
                }
                Err(e) => {
                    error!("{}", e);
                    exit(1);
                }
            },
            Some(device) => device,
            None => "default",
        }
        .to_string();

        info!("Using AlsaSink with format: {:?}", format);

        Self {
            pcm: None,
            format,
            device: name,
            period_buffer: vec![],
        }
    }
}

impl Sink for AlsaSink {
    fn start(&mut self) -> SinkResult<()> {
        if self.pcm.is_none() {
            let (pcm, bytes_per_period) = open_device(&self.device, self.format)?;
            self.pcm = Some(pcm);

            let current_capacity = self.period_buffer.capacity();

            if current_capacity > bytes_per_period {
                self.period_buffer.truncate(bytes_per_period);
                self.period_buffer.shrink_to_fit();
            } else if current_capacity < bytes_per_period {
                let extra = bytes_per_period - self.period_buffer.len();
                self.period_buffer.reserve_exact(extra);
            }

            // Should always match the "Period Buffer size in bytes: " trace! message.
            trace!(
                "Period Buffer capacity: {:?}",
                self.period_buffer.capacity()
            );
        }

        Ok(())
    }

    fn stop(&mut self) -> SinkResult<()> {
        // Zero fill the remainder of the period buffer and
        // write any leftover data before draining the actual PCM buffer.
        self.period_buffer.resize(self.period_buffer.capacity(), 0);
        self.write_buf()?;

        let pcm = self.pcm.as_mut().ok_or(AlsaError::NotConnected)?;

        pcm.drain().map_err(AlsaError::DrainFailure)?;

        self.pcm = None;
        Ok(())
    }

    sink_as_bytes!();
}

impl SinkAsBytes for AlsaSink {
    fn write_bytes(&mut self, data: &[u8]) -> SinkResult<()> {
        let mut start_index = 0;
        let data_len = data.len();
        let capacity = self.period_buffer.capacity();

        loop {
            let data_left = data_len - start_index;
            let space_left = capacity - self.period_buffer.len();
            let data_to_buffer = min(data_left, space_left);
            let end_index = start_index + data_to_buffer;

            self.period_buffer
                .extend_from_slice(&data[start_index..end_index]);

            if self.period_buffer.len() == capacity {
                self.write_buf()?;
            }

            if end_index == data_len {
                break Ok(());
            }

            start_index = end_index;
        }
    }
}

impl AlsaSink {
    pub const NAME: &'static str = "alsa";

    fn write_buf(&mut self) -> SinkResult<()> {
        let pcm = self.pcm.as_mut().ok_or(AlsaError::NotConnected)?;

        if let Err(e) = pcm.io_bytes().writei(&self.period_buffer) {
            // Capture and log the original error as a warning, and then try to recover.
            // If recovery fails then forward that error back to player.
            warn!(
                "Error writing from AlsaSink buffer to PCM, trying to recover, {}",
                e
            );

            pcm.try_recover(e, false).map_err(AlsaError::OnWrite)?
        }

        self.period_buffer.clear();
        Ok(())
    }
}
<commit_msg>Fix clippy comparison chain warning (#857)<commit_after>use super::{Open, Sink, SinkAsBytes, SinkError, SinkResult};
use crate::config::AudioFormat;
use crate::convert::Converter;
use crate::decoder::AudioPacket;
use crate::{NUM_CHANNELS, SAMPLE_RATE};
use alsa::device_name::HintIter;
use alsa::pcm::{Access, Format, HwParams, PCM};
use alsa::{Direction, ValueOr};
use std::cmp::min;
use std::cmp::Ordering;
use std::process::exit;
use std::time::Duration;
use thiserror::Error;

// 0.5 sec buffer.
const PERIOD_TIME: Duration = Duration::from_millis(100);
const BUFFER_TIME: Duration = Duration::from_millis(500);

#[derive(Debug, Error)]
enum AlsaError {
    #[error("<AlsaSink> Device {device} Unsupported Format {alsa_format:?} ({format:?}), {e}")]
    UnsupportedFormat {
        device: String,
        alsa_format: Format,
        format: AudioFormat,
        e: alsa::Error,
    },

    #[error("<AlsaSink> Device {device} Unsupported Channel Count {channel_count}, {e}")]
    UnsupportedChannelCount {
        device: String,
        channel_count: u8,
        e: alsa::Error,
    },

    #[error("<AlsaSink> Device {device} Unsupported Sample Rate {samplerate}, {e}")]
    UnsupportedSampleRate {
        device: String,
        samplerate: u32,
        e: alsa::Error,
    },

    #[error("<AlsaSink> Device {device} Unsupported Access Type RWInterleaved, {e}")]
    UnsupportedAccessType { device: String, e: alsa::Error },

    #[error("<AlsaSink> Device {device} May be Invalid, Busy, or Already in Use, {e}")]
    PcmSetUp { device: String, e: alsa::Error },

    #[error("<AlsaSink> Failed to Drain PCM Buffer, {0}")]
    DrainFailure(alsa::Error),

    #[error("<AlsaSink> {0}")]
    OnWrite(alsa::Error),

    #[error("<AlsaSink> Hardware, {0}")]
    HwParams(alsa::Error),

    #[error("<AlsaSink> Software, {0}")]
    SwParams(alsa::Error),

    #[error("<AlsaSink> PCM, {0}")]
    Pcm(alsa::Error),

    #[error("<AlsaSink> Could Not Parse Ouput Name(s) and/or Description(s)")]
    Parsing,

    #[error("<AlsaSink>")]
    NotConnected,
}

impl From<AlsaError> for SinkError {
    fn from(e: AlsaError) -> SinkError {
        use AlsaError::*;
        let es = e.to_string();
        match e {
            DrainFailure(_) | OnWrite(_) => SinkError::OnWrite(es),
            PcmSetUp { .. } => SinkError::ConnectionRefused(es),
            NotConnected => SinkError::NotConnected(es),
            _ => SinkError::InvalidParams(es),
        }
    }
}

pub struct AlsaSink {
    pcm: Option<PCM>,
    format: AudioFormat,
    device: String,
    period_buffer: Vec<u8>,
}

fn list_outputs() -> SinkResult<()> {
    println!("Listing available Alsa outputs:");
    for t in &["pcm", "ctl", "hwdep"] {
        println!("{} devices:", t);

        let i = HintIter::new_str(None, t).map_err(|_| AlsaError::Parsing)?;

        for a in i {
            if let Some(Direction::Playback) = a.direction {
                // mimic aplay -L
                let name = a.name.ok_or(AlsaError::Parsing)?;
                let desc = a.desc.ok_or(AlsaError::Parsing)?;

                println!("{}\n\t{}\n", name, desc.replace("\n", "\n\t"));
            }
        }
    }

    Ok(())
}

fn open_device(dev_name: &str, format: AudioFormat) -> SinkResult<(PCM, usize)> {
    let pcm = PCM::new(dev_name, Direction::Playback, false).map_err(|e| AlsaError::PcmSetUp {
        device: dev_name.to_string(),
        e,
    })?;

    let alsa_format = match format {
        AudioFormat::F64 => Format::float64(),
        AudioFormat::F32 => Format::float(),
        AudioFormat::S32 => Format::s32(),
        AudioFormat::S24 => Format::s24(),
        AudioFormat::S16 => Format::s16(),

        #[cfg(target_endian = "little")]
        AudioFormat::S24_3 => Format::S243LE,
        #[cfg(target_endian = "big")]
        AudioFormat::S24_3 => Format::S243BE,
    };

    let bytes_per_period = {
        let hwp = HwParams::any(&pcm).map_err(AlsaError::HwParams)?;

        hwp.set_access(Access::RWInterleaved)
            .map_err(|e| AlsaError::UnsupportedAccessType {
                device: dev_name.to_string(),
                e,
            })?;

        hwp.set_format(alsa_format)
            .map_err(|e| AlsaError::UnsupportedFormat {
                device: dev_name.to_string(),
                alsa_format,
                format,
                e,
            })?;

        hwp.set_rate(SAMPLE_RATE, ValueOr::Nearest).map_err(|e| {
            AlsaError::UnsupportedSampleRate {
                device: dev_name.to_string(),
                samplerate: SAMPLE_RATE,
                e,
            }
        })?;

        hwp.set_channels(NUM_CHANNELS as u32)
            .map_err(|e| AlsaError::UnsupportedChannelCount {
                device: dev_name.to_string(),
                channel_count: NUM_CHANNELS,
                e,
            })?;

        hwp.set_buffer_time_near(BUFFER_TIME.as_micros() as u32, ValueOr::Nearest)
            .map_err(AlsaError::HwParams)?;

        hwp.set_period_time_near(PERIOD_TIME.as_micros() as u32, ValueOr::Nearest)
            .map_err(AlsaError::HwParams)?;

        pcm.hw_params(&hwp).map_err(AlsaError::Pcm)?;

        let swp = pcm.sw_params_current().map_err(AlsaError::Pcm)?;

        // Don't assume we got what we wanted. Ask to make sure.
        let frames_per_period = hwp.get_period_size().map_err(AlsaError::HwParams)?;

        let frames_per_buffer = hwp.get_buffer_size().map_err(AlsaError::HwParams)?;

        swp.set_start_threshold(frames_per_buffer - frames_per_period)
            .map_err(AlsaError::SwParams)?;

        pcm.sw_params(&swp).map_err(AlsaError::Pcm)?;

        trace!("Frames per Buffer: {:?}", frames_per_buffer);
        trace!("Frames per Period: {:?}", frames_per_period);

        // Let ALSA do the math for us.
        pcm.frames_to_bytes(frames_per_period) as usize
    };

    trace!("Period Buffer size in bytes: {:?}", bytes_per_period);

    Ok((pcm, bytes_per_period))
}

impl Open for AlsaSink {
    fn open(device: Option<String>, format: AudioFormat) -> Self {
        let name = match device.as_deref() {
            Some("?") => match list_outputs() {
                Ok(_) => {
                    exit(0);
                }
                Err(e) => {
                    error!("{}", e);
                    exit(1);
                }
            },
            Some(device) => device,
            None => "default",
        }
        .to_string();

        info!("Using AlsaSink with format: {:?}", format);

        Self {
            pcm: None,
            format,
            device: name,
            period_buffer: vec![],
        }
    }
}

impl Sink for AlsaSink {
    fn start(&mut self) -> SinkResult<()> {
        if self.pcm.is_none() {
            let (pcm, bytes_per_period) = open_device(&self.device, self.format)?;
            self.pcm = Some(pcm);

            match self.period_buffer.capacity().cmp(&bytes_per_period) {
                Ordering::Greater => {
                    self.period_buffer.truncate(bytes_per_period);
                    self.period_buffer.shrink_to_fit();
                }
                Ordering::Less => {
                    let extra = bytes_per_period - self.period_buffer.len();
                    self.period_buffer.reserve_exact(extra);
                }
                Ordering::Equal => (),
            }

            // Should always match the "Period Buffer size in bytes: " trace! message.
            trace!(
                "Period Buffer capacity: {:?}",
                self.period_buffer.capacity()
            );
        }

        Ok(())
    }

    fn stop(&mut self) -> SinkResult<()> {
        // Zero fill the remainder of the period buffer and
        // write any leftover data before draining the actual PCM buffer.
        self.period_buffer.resize(self.period_buffer.capacity(), 0);
        self.write_buf()?;

        let pcm = self.pcm.take().ok_or(AlsaError::NotConnected)?;

        pcm.drain().map_err(AlsaError::DrainFailure)?;

        Ok(())
    }

    sink_as_bytes!();
}

impl SinkAsBytes for AlsaSink {
    fn write_bytes(&mut self, data: &[u8]) -> SinkResult<()> {
        let mut start_index = 0;
        let data_len = data.len();
        let capacity = self.period_buffer.capacity();

        loop {
            let data_left = data_len - start_index;
            let space_left = capacity - self.period_buffer.len();
            let data_to_buffer = min(data_left, space_left);
            let end_index = start_index + data_to_buffer;

            self.period_buffer
                .extend_from_slice(&data[start_index..end_index]);

            if self.period_buffer.len() == capacity {
                self.write_buf()?;
            }

            if end_index == data_len {
                break Ok(());
            }

            start_index = end_index;
        }
    }
}

impl AlsaSink {
    pub const NAME: &'static str = "alsa";

    fn write_buf(&mut self) -> SinkResult<()> {
        let pcm = self.pcm.as_mut().ok_or(AlsaError::NotConnected)?;

        if let Err(e) = pcm.io_bytes().writei(&self.period_buffer) {
            // Capture and log the original error as a warning, and then try to recover.
            // If recovery fails then forward that error back to player.
            warn!(
                "Error writing from AlsaSink buffer to PCM, trying to recover, {}",
                e
            );

            pcm.try_recover(e, false).map_err(AlsaError::OnWrite)?
        }

        self.period_buffer.clear();
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Contains the ThreadPool and the individual Threads.

use std::alloc::{Alloc, Layout, Global, handle_alloc_error};
use std::sync::atomic::{AtomicBool,Ordering};
use std::thread::{JoinHandle,self};
use std::sync::{Once, ONCE_INIT};
use std::ptr::NonNull;
use std::{ptr,mem};
use std::cell::UnsafeCell;

use pleco::MoveList;
use pleco::tools::pleco_arc::Arc;
use pleco::board::*;
use pleco::core::piece_move::BitMove;

use sync::LockLatch;
use time::uci_timer::*;
use search::Searcher;

use consts::*;

const KILOBYTE: usize = 1000;
const THREAD_STACK_SIZE: usize = 18000 * KILOBYTE;
const POOL_SIZE: usize = mem::size_of::<ThreadPool>();

// An object that is the same size as a thread pool.
type DummyThreadPool = [u8; POOL_SIZE];

// The Global threadpool! Yes, this is *technically* an array the same
// size as a ThreadPool object. This is a cheap hack to get a global value, as
// Rust isn't particularily fond of mutable global statics.
pub static mut THREADPOOL: DummyThreadPool = [0; POOL_SIZE];

// ONCE for the Threadpool
static THREADPOOL_INIT: Once = ONCE_INIT;

// Initializes the threadpool, called once on startup.
#[cold]
pub fn init_threadpool() {
    THREADPOOL_INIT.call_once(|| {
        unsafe {
            // We have a spawned thread create all structures, as a stack overflow can
            // occur otherwise
            let builder = thread::Builder::new()
                .name("Starter".to_string())
                .stack_size(THREAD_STACK_SIZE);

            let handle = builder.spawn_unchecked(
                move || {
                    let pool: *mut ThreadPool = mem::transmute(&mut THREADPOOL);
                    ptr::write(pool, ThreadPool::new());
            });
            handle.unwrap().join().unwrap();
        }
    });
}

/// Returns access to the global thread pool.
#[inline(always)]
pub fn threadpool() -> &'static mut ThreadPool {
    unsafe {
        mem::transmute::<&mut DummyThreadPool, &'static mut ThreadPool>(&mut THREADPOOL)
    }
}

#[derive(Copy, Clone)]
enum ThreadSelection {
    Main,
    NonMain,
    All
}

impl ThreadSelection {
    #[inline(always)]
    pub fn is_selection(self, id: usize) -> bool {
        match self {
            ThreadSelection::Main => id == 0,
            ThreadSelection::NonMain => id != 0,
            ThreadSelection::All => true
        }
    }
}

// Dummy struct to allow us to pass a pointer into a spawned thread.
struct SearcherPtr {
    ptr: UnsafeCell<*mut Searcher>
}

unsafe impl Sync for SearcherPtr {}
unsafe impl Send for SearcherPtr {}

/// The thread-pool for the chess engine.
pub struct ThreadPool {
    /// Access to each thread's Structure
    pub threads: Vec<UnsafeCell<*mut Searcher>>,
    /// Handles of each thread
    handles: Vec<JoinHandle<()>>,
    /// Condition for the main thread to start.
    pub main_cond: Arc<LockLatch>,
    /// Condition for all non-main threads
    pub thread_cond: Arc<LockLatch>,
    /// Stop condition, if true the threads should halt.
    pub stop: AtomicBool
}

// Okay, this all looks like madness, but there is some reason to it all.
// Basically, `ThreadPool` manages spawning and despawning threads, as well
// as passing state to / from those threads, telling them to stop, go, drop,
// and lastly determining the "best move" from all the threads.
///
// While we spawn all the other threads, We mostly communicate with the
// MainThread to do anything useful. The mainthread handles anything fun.
// The goal of the ThreadPool is to be NON BLOCKING, unless we want to await a
// result.
impl ThreadPool {

    /// Creates a new `ThreadPool`
    pub fn new() -> Self {
        let mut pool: ThreadPool = ThreadPool {
            threads: Vec::new(),
            handles: Vec::new(),
            main_cond: Arc::new(LockLatch::new()),
            thread_cond: Arc::new(LockLatch::new()),
            stop: AtomicBool::new(true)
        };
        // Lock both the cond variables
        pool.main_cond.lock();
        pool.thread_cond.lock();

        // spawn the main thread
        pool.attach_thread();
        pool
    }

    /// Spawns a new thread and appends it to our vector of JoinHandles.
    fn attach_thread(&mut self) {
         unsafe {
             let thread_ptr: SearcherPtr = self.create_thread();
             let builder = thread::Builder::new()
                 .name(self.size().to_string())
                 .stack_size(THREAD_STACK_SIZE);

             let handle = builder.spawn_unchecked(
                move || {
                    let thread = &mut **thread_ptr.ptr.get();
                    thread.cond.lock();
                    thread.idle_loop();
             }).unwrap();
             self.handles.push(handle);
        };
    }

    /// Allocates a thread structure and pushes it to the threadstack.
    ///
    /// This does not spawn a thread, just creates the structure in the heap for the thread.
    ///
    /// Only to be called by `attach_thread()`.
    fn create_thread(&mut self) -> SearcherPtr {
        let len: usize = self.threads.len();
        let layout = Layout::new::<Searcher>();
        let cond = if len == 0 {self.main_cond.clone()} else {self.thread_cond.clone()};
        unsafe {
            let result = Global.alloc_zeroed(layout);
            let new_ptr: *mut Searcher = match result {
                Ok(ptr) => ptr.cast().as_ptr() as *mut Searcher,
                Err(_err) => handle_alloc_error(layout),
            };
            ptr::write(new_ptr, Searcher::new(len, cond));
            self.threads.push(UnsafeCell::new(new_ptr));
            SearcherPtr {ptr: UnsafeCell::new(new_ptr)}
        }
    }

    /// Returns the number of threads
    #[inline(always)]
    pub fn size(&self) -> usize {
        self.threads.len()
    }

    /// Returns a pointer to the main thread.
    fn main(&mut self) -> &mut Searcher {
        unsafe {
            let main_thread: *mut Searcher = *self.threads.get_unchecked(0).get();
            &mut *main_thread
        }
    }

    /// Sets the use of standard out. This can be changed mid search as well.
    #[inline(always)]
    pub fn stdout(&mut self, use_stdout: bool) {
        USE_STDOUT.store(use_stdout, Ordering::Relaxed);
    }

    /// Sets the thread count of the pool. If num is less than 1, nothing will happen.
    ///
    /// # Safety
    ///
    /// Completely unsafe to use when the pool is searching.
    pub fn set_thread_count(&mut self, mut num: usize) {
        if num >= 1 {
            num = num.min(MAX_THREADS);
            self.wait_for_finish();
            self.kill_all();
            while self.size() < num {
                self.attach_thread();
            }
        }
    }

    /// Kills and de-allocates all the threads that are running. This function will also
    /// block on waiting for the search to finish.
    pub fn kill_all(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
        self.wait_for_finish();
        let mut join_handles = Vec::with_capacity(self.size());
        unsafe {
            // tell each thread to drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| { s.kill.store(true, Ordering::SeqCst) });

            // If awaiting a signal, wake up each thread so each can drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| { s.cond.set(); });

            // Start connecting each join handle. We don't unwrap here, as if one of the
            // threads fail, the other threads remain un-joined.
            while let Some(handle) = self.handles.pop() {
                join_handles.push(handle.join());
            }

            // De-allocate each thread.
            while let Some(unc) = self.threads.pop() {
                let th: *mut Searcher = *unc.get();
                let ptr: NonNull<u8> = mem::transmute(NonNull::new_unchecked(th));
                let layout = Layout::new::<Searcher>();
                Global.dealloc(ptr, layout);
            }
        }

        // Unwrap the results from each `thread::join`,
        while let Some(handle_result) = join_handles.pop() {
            handle_result.unwrap_or_else(|e| println!("Thread failed: {:?}",e));
        }
    }

    /// Sets the threads to stop (or not!).
    #[inline(always)]
    pub fn set_stop(&mut self, stop: bool) {
        self.stop.store(stop, Ordering::Relaxed);
    }

    /// Waits for all the threads to finish
    pub fn wait_for_finish(&self) {
        self.await_search_cond(ThreadSelection::All, false);
    }

    /// Waits for all the threads to start.
    pub fn wait_for_start(&self) {
        self.await_search_cond(ThreadSelection::All, true);
    }

    /// Waits for all non-main threads to finish.
    pub fn wait_for_non_main(&self) {
        self.await_search_cond(ThreadSelection::NonMain, true);
    }

    /// Waits for all the non-main threads to start running
    pub fn wait_for_main_start(&self) {
        self.await_search_cond(ThreadSelection::Main, true);
    }

    fn await_search_cond(&self, thread_sel: ThreadSelection, await_search: bool) {
        self.threads.iter()
            .map(|s| unsafe {&**s.get()})
            .filter(|t| thread_sel.is_selection(t.id))
            .for_each(|t: &Searcher|{ t.searching.await(await_search); });
    }

    pub fn clear_all(&mut self) {
        self.threads.iter_mut()
            .map(|thread_ptr| unsafe { &mut **(*thread_ptr).get() })
            .for_each(|t| t.clear());
    }

    /// Starts a UCI search. The result will be printed to stdout if the stdout setting
    /// is true.
    pub fn uci_search(&mut self, board: &Board, limits: &Limits) {

        // Start the timer!
        if let Some(uci_timer) = limits.use_time_management() {
            timer().init(limits.start, &uci_timer, board.turn(), board.moves_played());
        } else {
            timer().start_timer(limits.start);
        }

        let root_moves: MoveList = board.generate_moves();

        assert!(!root_moves.is_empty());
        self.wait_for_finish();
        self.stop.store(false, Ordering::Relaxed);

        for thread_ptr in self.threads.iter_mut() {
            let thread: &mut Searcher = unsafe {&mut **(*thread_ptr).get()};
            thread.nodes.store(0, Ordering::Relaxed);
            thread.depth_completed = 0;
            thread.board = board.shallow_clone();
            thread.limit = limits.clone();
            thread.root_moves().replace(&root_moves);
        }

        self.main_cond.set();
        self.wait_for_main_start();
        self.main_cond.lock();
    }

    /// Performs a standard search, and blocks waiting for a returned `BitMove`.
    pub fn search(&mut self, board: &Board, limits: &Limits) -> BitMove {
        self.uci_search(board, limits);
        self.wait_for_finish();
        self.best_move()
    }

    /// Returns the best move of a search
    pub fn best_move(&mut self) -> BitMove {
        self.main().root_moves().get(0).unwrap().bit_move
    }

    /// Returns total number of nodes searched so far.
    pub fn nodes(&self) -> u64 {
        self.threads.iter()
            .map(|s| unsafe {&**s.get()})
            .map(|s: &Searcher| s.nodes.load(Ordering::Relaxed))
            .sum()
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        self.kill_all();
    }
}
<commit_msg>Bug fix<commit_after>//! Contains the ThreadPool and the individual Threads.

use std::alloc::{Alloc, Layout, Global, handle_alloc_error};
use std::sync::atomic::{AtomicBool,Ordering};
use std::thread::{JoinHandle,self};
use std::sync::{Once, ONCE_INIT};
use std::ptr::NonNull;
use std::{ptr,mem};
use std::cell::UnsafeCell;

use pleco::MoveList;
use pleco::tools::pleco_arc::Arc;
use pleco::board::*;
use pleco::core::piece_move::BitMove;

use sync::LockLatch;
use time::uci_timer::*;
use search::Searcher;

use consts::*;

const KILOBYTE: usize = 1000;
const THREAD_STACK_SIZE: usize = 18000 * KILOBYTE;
const POOL_SIZE: usize = mem::size_of::<ThreadPool>();

// An object that is the same size as a thread pool.
type DummyThreadPool = [u8; POOL_SIZE];

// The Global threadpool! Yes, this is *technically* an array the same
// size as a ThreadPool object. This is a cheap hack to get a global value, as
// Rust isn't particularily fond of mutable global statics.
pub static mut THREADPOOL: DummyThreadPool = [0; POOL_SIZE];

// ONCE for the Threadpool
static THREADPOOL_INIT: Once = ONCE_INIT;

// Initializes the threadpool, called once on startup.
#[cold]
pub fn init_threadpool() {
    THREADPOOL_INIT.call_once(|| {
        unsafe {
            // We have a spawned thread create all structures, as a stack overflow can
            // occur otherwise
            let builder = thread::Builder::new()
                .name("Starter".to_string())
                .stack_size(THREAD_STACK_SIZE);

            let handle = builder.spawn_unchecked(
                move || {
                    let pool: *mut ThreadPool = mem::transmute(&mut THREADPOOL);
                    ptr::write(pool, ThreadPool::new());
            });
            handle.unwrap().join().unwrap();
        }
    });
}

/// Returns access to the global thread pool.
#[inline(always)]
pub fn threadpool() -> &'static mut ThreadPool {
    unsafe {
        mem::transmute::<&mut DummyThreadPool, &'static mut ThreadPool>(&mut THREADPOOL)
    }
}

#[derive(Copy, Clone)]
enum ThreadSelection {
    Main,
    NonMain,
    All
}

impl ThreadSelection {
    #[inline(always)]
    pub fn is_selection(self, id: usize) -> bool {
        match self {
            ThreadSelection::Main => id == 0,
            ThreadSelection::NonMain => id != 0,
            ThreadSelection::All => true
        }
    }
}

// Dummy struct to allow us to pass a pointer into a spawned thread.
struct SearcherPtr {
    ptr: UnsafeCell<*mut Searcher>
}

unsafe impl Sync for SearcherPtr {}
unsafe impl Send for SearcherPtr {}

/// The thread-pool for the chess engine.
pub struct ThreadPool {
    /// Access to each thread's Structure
    pub threads: Vec<UnsafeCell<*mut Searcher>>,
    /// Handles of each thread
    handles: Vec<JoinHandle<()>>,
    /// Condition for the main thread to start.
    pub main_cond: Arc<LockLatch>,
    /// Condition for all non-main threads
    pub thread_cond: Arc<LockLatch>,
    /// Stop condition, if true the threads should halt.
    pub stop: AtomicBool
}

// Okay, this all looks like madness, but there is some reason to it all.
// Basically, `ThreadPool` manages spawning and despawning threads, as well
// as passing state to / from those threads, telling them to stop, go, drop,
// and lastly determining the "best move" from all the threads.
///
// While we spawn all the other threads, We mostly communicate with the
// MainThread to do anything useful. The mainthread handles anything fun.
// The goal of the ThreadPool is to be NON BLOCKING, unless we want to await a
// result.
impl ThreadPool {

    /// Creates a new `ThreadPool`
    pub fn new() -> Self {
        let mut pool: ThreadPool = ThreadPool {
            threads: Vec::new(),
            handles: Vec::new(),
            main_cond: Arc::new(LockLatch::new()),
            thread_cond: Arc::new(LockLatch::new()),
            stop: AtomicBool::new(true)
        };
        // Lock both the cond variables
        pool.main_cond.lock();
        pool.thread_cond.lock();

        // spawn the main thread
        pool.attach_thread();
        pool
    }

    /// Spawns a new thread and appends it to our vector of JoinHandles.
    fn attach_thread(&mut self) {
         unsafe {
             let thread_ptr: SearcherPtr = self.create_thread();
             let builder = thread::Builder::new()
                 .name(self.size().to_string())
                 .stack_size(THREAD_STACK_SIZE);

             let handle = builder.spawn_unchecked(
                move || {
                    let thread = &mut **thread_ptr.ptr.get();
                    thread.cond.lock();
                    thread.idle_loop();
             }).unwrap();
             self.handles.push(handle);
        };
    }

    /// Allocates a thread structure and pushes it to the threadstack.
    ///
    /// This does not spawn a thread, just creates the structure in the heap for the thread.
    ///
    /// Only to be called by `attach_thread()`.
    fn create_thread(&mut self) -> SearcherPtr {
        let len: usize = self.threads.len();
        let layout = Layout::new::<Searcher>();
        let cond = if len == 0 {self.main_cond.clone()} else {self.thread_cond.clone()};
        unsafe {
            let result = Global.alloc_zeroed(layout);
            let new_ptr: *mut Searcher = match result {
                Ok(ptr) => ptr.cast().as_ptr() as *mut Searcher,
                Err(_err) => handle_alloc_error(layout),
            };
            ptr::write(new_ptr, Searcher::new(len, cond));
            self.threads.push(UnsafeCell::new(new_ptr));
            SearcherPtr {ptr: UnsafeCell::new(new_ptr)}
        }
    }

    /// Returns the number of threads
    #[inline(always)]
    pub fn size(&self) -> usize {
        self.threads.len()
    }

    /// Returns a pointer to the main thread.
    fn main(&mut self) -> &mut Searcher {
        unsafe {
            let main_thread: *mut Searcher = *self.threads.get_unchecked(0).get();
            &mut *main_thread
        }
    }

    /// Sets the use of standard out. This can be changed mid search as well.
    #[inline(always)]
    pub fn stdout(&mut self, use_stdout: bool) {
        USE_STDOUT.store(use_stdout, Ordering::Relaxed);
    }

    /// Sets the thread count of the pool. If num is less than 1, nothing will happen.
    ///
    /// # Safety
    ///
    /// Completely unsafe to use when the pool is searching.
    pub fn set_thread_count(&mut self, mut num: usize) {
        if num >= 1 {
            num = num.min(MAX_THREADS);
            self.wait_for_finish();
            self.kill_all();
            while self.size() < num {
                self.attach_thread();
            }
        }
    }

    /// Kills and de-allocates all the threads that are running. This function will also
    /// block on waiting for the search to finish.
    pub fn kill_all(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
        self.wait_for_finish();
        let mut join_handles = Vec::with_capacity(self.size());
        unsafe {
            // tell each thread to drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| { s.kill.store(true, Ordering::SeqCst) });

            // If awaiting a signal, wake up each thread so each can drop
            self.threads.iter()
                .map(|s| &**s.get())
                .for_each(|s: &Searcher| { s.cond.set(); });

            // Start connecting each join handle. We don't unwrap here, as if one of the
            // threads fail, the other threads remain un-joined.
            while let Some(handle) = self.handles.pop() {
                join_handles.push(handle.join());
            }

            // De-allocate each thread.
            while let Some(unc) = self.threads.pop() {
                let th: *mut Searcher = *unc.get();
                let ptr: NonNull<u8> = mem::transmute(NonNull::new_unchecked(th));
                let layout = Layout::new::<Searcher>();
                Global.dealloc(ptr, layout);
            }
        }

        // Unwrap the results from each `thread::join`,
        while let Some(handle_result) = join_handles.pop() {
            handle_result.unwrap_or_else(|e| println!("Thread failed: {:?}",e));
        }
    }

    /// Sets the threads to stop (or not!).
    #[inline(always)]
    pub fn set_stop(&mut self, stop: bool) {
        self.stop.store(stop, Ordering::Relaxed);
    }

    /// Waits for all the threads to finish
    pub fn wait_for_finish(&self) {
        self.await_search_cond(ThreadSelection::All, false);
    }

    /// Waits for all the threads to start.
    pub fn wait_for_start(&self) {
        self.await_search_cond(ThreadSelection::All, true);
    }

    /// Waits for all non-main threads to finish.
    pub fn wait_for_non_main(&self) {
        self.await_search_cond(ThreadSelection::NonMain, false);
    }

    /// Waits for all the non-main threads to start running
    pub fn wait_for_main_start(&self) {
        self.await_search_cond(ThreadSelection::Main, true);
    }

    fn await_search_cond(&self, thread_sel: ThreadSelection, await_search: bool) {
        self.threads.iter()
            .map(|s| unsafe {&**s.get()})
            .filter(|t| thread_sel.is_selection(t.id))
            .for_each(|t: &Searcher|{ t.searching.await(await_search); });
    }

    pub fn clear_all(&mut self) {
        self.threads.iter_mut()
            .map(|thread_ptr| unsafe { &mut **(*thread_ptr).get() })
            .for_each(|t| t.clear());
    }

    /// Starts a UCI search. The result will be printed to stdout if the stdout setting
    /// is true.
    pub fn uci_search(&mut self, board: &Board, limits: &Limits) {

        // Start the timer!
        if let Some(uci_timer) = limits.use_time_management() {
            timer().init(limits.start, &uci_timer, board.turn(), board.moves_played());
        } else {
            timer().start_timer(limits.start);
        }

        let root_moves: MoveList = board.generate_moves();

        assert!(!root_moves.is_empty());
        self.wait_for_finish();
        self.stop.store(false, Ordering::Relaxed);

        for thread_ptr in self.threads.iter_mut() {
            let thread: &mut Searcher = unsafe {&mut **(*thread_ptr).get()};
            thread.nodes.store(0, Ordering::Relaxed);
            thread.depth_completed = 0;
            thread.board = board.shallow_clone();
            thread.limit = limits.clone();
            thread.root_moves().replace(&root_moves);
        }

        self.main_cond.set();
        self.wait_for_main_start();
        self.main_cond.lock();
    }

    /// Performs a standard search, and blocks waiting for a returned `BitMove`.
    pub fn search(&mut self, board: &Board, limits: &Limits) -> BitMove {
        self.uci_search(board, limits);
        self.wait_for_finish();
        self.best_move()
    }

    /// Returns the best move of a search
    pub fn best_move(&mut self) -> BitMove {
        self.main().root_moves().get(0).unwrap().bit_move
    }

    /// Returns total number of nodes searched so far.
    pub fn nodes(&self) -> u64 {
        self.threads.iter()
            .map(|s| unsafe {&**s.get()})
            .map(|s: &Searcher| s.nodes.load(Ordering::Relaxed))
            .sum()
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        self.kill_all();
    }
}
<|endoftext|>


<commit_before>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::ops::{Div, Sub};
use std::path::{Path, PathBuf};

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::{MAGIC_CODE_NIP1, MAGIC_CODE_NI1},
    util::{is_gz_file, is_hdr_file},
    volume::element::DataElement,
    NiftiHeader, NiftiType, Result,
};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) = prepare_header_and_paths(&header_path, data, reference, A::DATA_TYPE);

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let header_file = File::create(&header_path)?;
    let mut header_writer = BufWriter::new(header_file);
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut header_writer, &header, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        let mut data_writer = BufWriter::new(data_file);
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_writer, compression_level);
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut data_writer, &header, data)?;
        }
    }

    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0.  If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    let is_gz = is_gz_file(&header_path);
    let (mut header, _) = prepare_header_and_paths(&header_path, data, reference, NiftiType::Rgb24);

    // The `scl_slope` and `scl_inter` fields are ignored on the Rgb24 type.
    header.scl_slope = 1.0;
    header.scl_inter = 0.0;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let f = File::create(&header_path)?;
    let mut writer = BufWriter::new(f);
    if is_gz {
        let mut e = GzEncoder::new(writer, Compression::fast());
        write_header(&mut e, &header)?;
        write_slices(&mut e, data)?;
        let _ = e.finish()?; // Must use result
    } else {
        write_header(&mut writer, &header)?;
        write_slices(&mut writer, data)?;
    }
    Ok(())
}

fn prepare_header_and_paths<P, T, D>(
    header_path: P,
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> (NiftiHeader, PathBuf)
where
    P: AsRef<Path>,
    T: Data,
    D: Dimension,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    let mut header = NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    };

    let mut path_buf = PathBuf::from(header_path.as_ref());
    let data_path = if is_hdr_file(&header_path) {
        header.vox_offset = 0.0;
        header.magic = *MAGIC_CODE_NI1;
        let data_path = if is_gz_file(&header_path) {
            let _ = path_buf.set_extension("");
            path_buf.with_extension("img.gz")
        } else {
            path_buf.with_extension("img")
        };
        data_path
    } else {
        path_buf
    };

    (header, data_path)
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
where
    W: WriteBytesExt,
{
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<T, D, W>(writer: &mut W, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
    T: Clone + PodTransmutable,
    T: Div<Output = T>,
    T: FromPrimitive,
    T: PodTransmutable,
    T: ScalarOperand,
    T: Sub<Output = T>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        // TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
        // voxel values to 0 if we pass an ndarray of integers.
        let slope = T::from_f32(slope).unwrap();
        let inter = T::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(writer, arr_data.sub(inter).div(slope))?;
        }
    } else {
        write_slices(writer, data)?;
    }
    Ok(())
}

fn write_slices<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice(writer, slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice(writer, slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: WriteBytesExt,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
    Ok(())
}
<commit_msg>Write RGB hdr<commit_after>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::ops::{Div, Sub};
use std::path::{Path, PathBuf};

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::{MAGIC_CODE_NI1, MAGIC_CODE_NIP1},
    util::{is_gz_file, is_hdr_file},
    volume::element::DataElement,
    NiftiHeader, NiftiType, Result,
};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) = prepare_header_and_paths(&header_path, data, reference, A::DATA_TYPE);

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    let mut header_writer = BufWriter::new(header_file);
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut header_writer, &header, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        let mut data_writer = BufWriter::new(data_file);
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_writer, compression_level);
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut data_writer, &header, data)?;
        }
    }

    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0.  If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (mut header, data_path) =
        prepare_header_and_paths(&header_path, data, reference, NiftiType::Rgb24);

    // The `scl_slope` and `scl_inter` fields are ignored on the Rgb24 type.
    header.scl_slope = 1.0;
    header.scl_inter = 0.0;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    let mut header_writer = BufWriter::new(header_file);
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            write_slices(&mut e, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_slices(&mut header_writer, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        let mut data_writer = BufWriter::new(data_file);
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_writer, compression_level);
            write_slices(&mut e, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_slices(&mut data_writer, data)?;
        }
    }

    Ok(())
}

fn prepare_header_and_paths<P, T, D>(
    header_path: P,
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> (NiftiHeader, PathBuf)
where
    P: AsRef<Path>,
    T: Data,
    D: Dimension,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    let mut header = NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    };

    let mut path_buf = PathBuf::from(header_path.as_ref());
    let data_path = if is_hdr_file(&header_path) {
        header.vox_offset = 0.0;
        header.magic = *MAGIC_CODE_NI1;
        let data_path = if is_gz_file(&header_path) {
            let _ = path_buf.set_extension("");
            path_buf.with_extension("img.gz")
        } else {
            path_buf.with_extension("img")
        };
        data_path
    } else {
        path_buf
    };

    (header, data_path)
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
where
    W: WriteBytesExt,
{
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<T, D, W>(writer: &mut W, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
    T: Clone + PodTransmutable,
    T: Div<Output = T>,
    T: FromPrimitive,
    T: PodTransmutable,
    T: ScalarOperand,
    T: Sub<Output = T>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        // TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
        // voxel values to 0 if we pass an ndarray of integers.
        let slope = T::from_f32(slope).unwrap();
        let inter = T::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(writer, arr_data.sub(inter).div(slope))?;
        }
    } else {
        write_slices(writer, data)?;
    }
    Ok(())
}

fn write_slices<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice(writer, slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice(writer, slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: WriteBytesExt,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
    Ok(())
}
<|endoftext|>