<commit_before>use dicom::core::smallvec;
use dicom::{core::dicom_value, dictionary_std::tags};
use dicom::{
    core::{DataElement, PrimitiveValue, VR},
    encoding::transfer_syntax,
    object::{mem::InMemDicomObject, open_file, StandardDataDictionary},
    transfer_syntax::TransferSyntaxRegistry,
};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use smallvec::smallvec;
use std::io::Write;
use std::path::PathBuf;
use structopt::StructOpt;
use transfer_syntax::TransferSyntaxIndex;

/// DICOM C-STORE SCU
#[derive(Debug, StructOpt)]
struct App {
    /// socket address to STORE SCP (example: "127.0.0.1:104")
    addr: String,
    /// the DICOM file to store
    file: PathBuf,
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// the C-STORE message ID
    #[structopt(short = "m", long = "message-id", default_value = "1")]
    message_id: u16,
    /// the calling AE title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[structopt(long = "called-ae-title", default_value = "ANY-SCP")]
    called_ae_title: String,
    /// the maximum PDU length accepted by the SCU
    #[structopt(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let App {
        addr,
        file,
        verbose,
        message_id,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
    } = App::from_args();

    if verbose {
        println!("Opening file '{}'...", file.display());
    }

    let dicom_file = open_file(file)?;

    let meta = dicom_file.meta();

    let storage_sop_class_uid = &meta.media_storage_sop_class_uid;
    let storage_sop_instance_uid = &meta.media_storage_sop_instance_uid;
    let transfer_syntax = &meta.transfer_syntax.trim_end_matches('\0');

    let file_ts_desc = TransferSyntaxRegistry
        .get(transfer_syntax)
        .ok_or("Unsupported file transfer syntax")?;

    if verbose {
        println!("Establishing association with '{}'...", &addr);
    }

    let mut scu = ClientAssociationOptions::new()
        .with_abstract_syntax(storage_sop_class_uid.to_string())
        .with_transfer_syntax(transfer_syntax.to_string())
        .calling_ae_title(calling_ae_title)
        .called_ae_title(called_ae_title)
        .max_pdu_length(max_pdu_length)
        .establish(addr)?;

    if verbose {
        println!("Association established");
    }

    // TODO(#106) transfer syntax conversion is currently not supported
    let pc_selected = if let Some(pc) = scu.presentation_contexts().iter().find(|pc| {
        // Check support for this transfer syntax.
        // If it is the same as the file, we're good.
        // Otherwise, uncompressed data set encoding
        // and native pixel data is required on both ends.
        let ts = &pc.transfer_syntax;
        ts == transfer_syntax
            || TransferSyntaxRegistry
                .get(&pc.transfer_syntax)
                .filter(|ts| file_ts_desc.is_codec_free() && ts.is_codec_free())
                .map(|_| true)
                .unwrap_or(false)
    }) {
        pc.clone()
    } else {
        eprintln!("Could not choose a transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    let ts = TransferSyntaxRegistry
        .get(&pc_selected.transfer_syntax)
        .ok_or("Poorly negotiated transfer syntax")?;

    if verbose {
        println!("Transfer Syntax: {}", ts.name());
    }

    let cmd = store_req_command(
        storage_sop_class_uid,
        storage_sop_instance_uid,
        message_id,
    );

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(
        &mut cmd_data,
        &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
    )?;

    let mut object_data = Vec::with_capacity(2048);
    dicom_file.write_dataset_with_ts(&mut object_data, ts)?;

    let nbytes = cmd_data.len() + object_data.len();

    if verbose {
        println!("Sending payload (~ {} kB)...", nbytes / 1_000);
    }

    if nbytes < scu.acceptor_max_pdu_length() as usize - 100 {
        let pdu = Pdu::PData {
            data: vec![
                PDataValue {
                    presentation_context_id: pc_selected.id,
                    value_type: PDataValueType::Command,
                    is_last: true,
                    data: cmd_data,
                },
                PDataValue {
                    presentation_context_id: pc_selected.id,
                    value_type: PDataValueType::Data,
                    is_last: true,
                    data: object_data,
                },
            ],
        };

        scu.send(&pdu)?;
    } else {
        let pdu = Pdu::PData {
            data: vec![PDataValue {
                presentation_context_id: pc_selected.id,
                value_type: PDataValueType::Command,
                is_last: true,
                data: cmd_data,
            }],
        };

        scu.send(&pdu)?;

        {
            let mut pdata = scu.send_pdata(pc_selected.id);
            pdata.write_all(&object_data)?;
        }
    }

    if verbose {
        println!("Awaiting response...");
    }

    let rsp_pdu = scu.receive()?;

    match rsp_pdu {
        Pdu::PData { data } => {
            let data_value = &data[0];

            let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                &data_value.data[..],
                &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
            )?;
            if verbose {
                println!("Response: {:?}", cmd_obj);
            }
            let status = cmd_obj.element(tags::STATUS)?.to_int::<u16>()?;
            let storage_sop_instance_uid =
                storage_sop_instance_uid.trim_end_matches(|c: char| c.is_whitespace() || c == '\0');
            if status == 0 {
                println!("Sucessfully stored instance `{}`", storage_sop_instance_uid);
            } else {
                println!(
                    "Failed to store instance `{}` (status code {})",
                    storage_sop_instance_uid, status
                );
            }

            scu.release()?;
        }

        pdu @ Pdu::Unknown { .. }
        | pdu @ Pdu::AssociationRQ { .. }
        | pdu @ Pdu::AssociationAC { .. }
        | pdu @ Pdu::AssociationRJ { .. }
        | pdu @ Pdu::ReleaseRQ
        | pdu @ Pdu::ReleaseRP
        | pdu @ Pdu::AbortRQ { .. } => {
            eprintln!("Unexpected SCP response: {:?}", pdu);
            let _ = scu.abort();
            std::process::exit(-2);
        }
    }

    Ok(())
}

fn store_req_command(
    storage_sop_class_uid: &str,
    storage_sop_instance_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::create_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            12 + 8
                + even_len(storage_sop_class_uid.len())
                + 10
                + 10
                + 10
                + 10
                + 8
                + even_len(storage_sop_instance_uid.len()),
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // affected SOP Instance UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_instance_uid),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<commit_msg>[storescu] Enable sending more than one file<commit_after>use dicom::core::smallvec;
use dicom::{core::dicom_value, dictionary_std::tags};
use dicom::{
    core::{DataElement, PrimitiveValue, VR},
    encoding::transfer_syntax,
    object::{mem::InMemDicomObject, open_file, StandardDataDictionary},
    transfer_syntax::TransferSyntaxRegistry,
};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use smallvec::smallvec;
use std::ffi::OsStr;
use std::io::Write;
use std::path::PathBuf;
use structopt::StructOpt;
use transfer_syntax::TransferSyntaxIndex;
use std::collections::HashSet;

/// DICOM C-STORE SCU
#[derive(Debug, StructOpt)]
struct App {
    /// socket address to STORE SCP (example: "127.0.0.1:104")
    addr: String,
    /// the DICOM file(s) to store
    files: Vec<PathBuf>,
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// the C-STORE message ID
    #[structopt(short = "m", long = "message-id", default_value = "1")]
    message_id: u16,
    /// the calling AE title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[structopt(long = "called-ae-title", default_value = "ANY-SCP")]
    called_ae_title: String,
    /// the maximum PDU length accepted by the SCU
    #[structopt(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
    /// fail if not all DICOM files can be transferred
    #[structopt(long = "fail-first")]
    fail_first: bool,
}

struct DicomFile {
    /// File path
    file: PathBuf,
    /// Storage SOP Class UID
    sop_class_uid: String,
    /// Storage SOP Instance UID
    sop_instance_uid: String,
    /// File Transfer Syntax
    file_transfer_syntax: String,
    /// Transfer Syntax selected
    ts_selected: Option<String>,
    /// Presentation Context selected
    pc_selected: Option<dicom_ul::pdu::PresentationContextResult>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let App {
        addr,
        files,
        verbose,
        message_id,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
        fail_first,
    } = App::from_args();

    let mut dicom_files: Vec<DicomFile> = vec![];
    let mut presentation_contexts = HashSet::new();

    for file in files {

        if verbose {
            println!("Opening file '{}'...", file.display());
        }

        match check_file(file) {
            Ok(dicom_file) => {
                presentation_contexts.insert((dicom_file.sop_class_uid.to_string(),dicom_file.file_transfer_syntax.clone()));
                dicom_files.push(dicom_file);
            },
            Err(e) => {
                if verbose {
                    println!("Error: {}", e);
                }
            }
        }
    }

    if dicom_files.is_empty() {
        eprintln!("No supported files to transfer");
        std::process::exit(-1);
    }

    if verbose {
        println!("Establishing association with '{}'...", &addr);
    }

    let mut scu_init = ClientAssociationOptions::new();

    for (storage_sop_class_uid,transfer_syntax) in &presentation_contexts {
        scu_init = scu_init.with_abstract_syntax(storage_sop_class_uid)
                    .with_transfer_syntax(transfer_syntax);
    }
    let mut scu = scu_init.calling_ae_title(calling_ae_title)
                    .called_ae_title(called_ae_title)
                    .max_pdu_length(max_pdu_length)
                    .establish(addr)?;

    if verbose {
        println!("Association established");
    }

    for mut file in &mut dicom_files {
        // TODO(#106) transfer syntax conversion is currently not supported
        match check_presentation_contexts(file, scu.presentation_contexts()) {
            Ok((pc, ts)) => {
                file.pc_selected = Some(pc);
                file.ts_selected = Some(ts);
            },
            Err(e) => {
                if fail_first {
                    eprintln!("Could not choose a transfer syntax: {}", e);
                    let _ = scu.abort();
                    std::process::exit(-2);
                }
            }
        }
    }

    for file in dicom_files {

        if let (Some(pc_selected),Some(ts_selected)) = (file.pc_selected, file.ts_selected) {

            let cmd = store_req_command(
                &file.sop_class_uid,
                &file.sop_instance_uid,
                message_id,
            );

            let mut cmd_data = Vec::with_capacity(128);
            cmd.write_dataset_with_ts(
                &mut cmd_data,
                &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
            )?;

            let mut object_data = Vec::with_capacity(2048);
            let dicom_file = open_file(file.file)?;
            let ts_selected = TransferSyntaxRegistry.get(&ts_selected).ok_or("Unsupported file transfer syntax")?;
            dicom_file.write_dataset_with_ts(&mut object_data, ts_selected)?;

            let nbytes = cmd_data.len() + object_data.len();

            if verbose {
                println!("Sending payload (~ {} kB)...", nbytes / 1_000);
            }

            if nbytes < scu.acceptor_max_pdu_length() as usize - 100 {
                let pdu = Pdu::PData {
                    data: vec![
                        PDataValue {
                            presentation_context_id: pc_selected.id,
                            value_type: PDataValueType::Command,
                            is_last: true,
                            data: cmd_data,
                        },
                        PDataValue {
                            presentation_context_id: pc_selected.id,
                            value_type: PDataValueType::Data,
                            is_last: true,
                            data: object_data,
                        },
                    ],
                };

                scu.send(&pdu)?;
            } else {
                let pdu = Pdu::PData {
                    data: vec![PDataValue {
                        presentation_context_id: pc_selected.id,
                        value_type: PDataValueType::Command,
                        is_last: true,
                        data: cmd_data,
                    }],
                };

                scu.send(&pdu)?;

                {
                    let mut pdata = scu.send_pdata(pc_selected.id);
                    pdata.write_all(&object_data)?;
                }
            }

            if verbose {
                println!("Awaiting response...");
            }

            let rsp_pdu = scu.receive()?;

            match rsp_pdu {
                Pdu::PData { data } => {
                    let data_value = &data[0];

                    let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                        &data_value.data[..],
                        &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
                    )?;
                    if verbose {
                        println!("Response: {:?}", cmd_obj);
                    }
                    let status = cmd_obj.element(tags::STATUS)?.to_int::<u16>()?;
                    let storage_sop_instance_uid =
                        file.sop_instance_uid.trim_end_matches(|c: char| c.is_whitespace() || c == '\0');
                    if status == 0 {
                        println!("Sucessfully stored instance `{}`", storage_sop_instance_uid);
                    } else {
                        println!(
                            "Failed to store instance `{}` (status code {})",
                            storage_sop_instance_uid, status
                        );
                        if fail_first {
                            let _ = scu.abort();
                            std::process::exit(-2);
                        }
                    }
                }

                pdu @ Pdu::Unknown { .. }
                | pdu @ Pdu::AssociationRQ { .. }
                | pdu @ Pdu::AssociationAC { .. }
                | pdu @ Pdu::AssociationRJ { .. }
                | pdu @ Pdu::ReleaseRQ
                | pdu @ Pdu::ReleaseRP
                | pdu @ Pdu::AbortRQ { .. } => {
                    eprintln!("Unexpected SCP response: {:?}", pdu);
                    let _ = scu.abort();
                    std::process::exit(-2);
                }
            }
        }
    }

scu.release()?;
    Ok(())
}

fn store_req_command(
    storage_sop_class_uid: &str,
    storage_sop_instance_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::create_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            12 + 8
                + even_len(storage_sop_class_uid.len())
                + 10
                + 10
                + 10
                + 10
                + 8
                + even_len(storage_sop_instance_uid.len()),
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // affected SOP Instance UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_instance_uid),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

fn check_file(file: PathBuf) -> Result<DicomFile, Box<dyn std::error::Error>> {
    // Ignore DICOMDIR files until better support is added
    let _ = (file.file_name() != Some(OsStr::new("DICOMDIR"))).then(|| false).ok_or("DICOMDIR file not supported")?;
    let dicom_file = open_file(&file)?;

    let meta = dicom_file.meta();

    let storage_sop_class_uid = &meta.media_storage_sop_class_uid;
    let storage_sop_instance_uid = &meta.media_storage_sop_instance_uid;
    let transfer_syntax_uid = &meta.transfer_syntax.trim_end_matches('\0');
    let ts = TransferSyntaxRegistry.get(transfer_syntax_uid).ok_or("Unsupported file transfer syntax")?;
    Ok(DicomFile{file,
                 sop_class_uid: storage_sop_class_uid.to_string(),
                 sop_instance_uid: storage_sop_instance_uid.to_string(),
                 file_transfer_syntax: String::from(ts.uid()),
                 ts_selected: None,
                 pc_selected: None,})
}

fn check_presentation_contexts(file: &DicomFile, pcs: &[dicom_ul::pdu::PresentationContextResult]) -> Result<(dicom_ul::pdu::PresentationContextResult, String), Box<dyn std::error::Error>> {
    let file_ts = TransferSyntaxRegistry.get(&file.file_transfer_syntax).ok_or("Unsupported file transfer syntax")?;
    // TODO(#106) transfer syntax conversion is currently not supported
    let pc = pcs.iter().find(|pc| {
        // Check support for this transfer syntax.
        // If it is the same as the file, we're good.
        // Otherwise, uncompressed data set encoding
        // and native pixel data is required on both ends.
        let ts = &pc.transfer_syntax;
        ts == file_ts.uid()
            || TransferSyntaxRegistry
               .get(&pc.transfer_syntax)
               .filter(|ts| file_ts.is_codec_free() && ts.is_codec_free())
               .map(|_| true)
               .unwrap_or(false)
    }).ok_or("No presentation context accepted")?;
    let ts = TransferSyntaxRegistry
             .get(&pc.transfer_syntax)
             .ok_or("Poorly negotiated transfer syntax")?;

    Ok((pc.clone(), String::from(ts.uid())))
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86` and `x86_64` descriptor tables (IDT, GDT, or LDT)

use memory::PAddr;

/// A pointer to a descriptor table.
/// This is a format suitable
#[repr(C, packed)]
pub struct Pointer { /// the length of the descriptor table
                     pub limit: u16
                   , /// pointer to the region in memory
                     /// containing the descriptor table.
                     pub base: PAddr
                   }

/// A descriptor table (IDT or GDT)
pub trait DTable {
    fn get_ptr(&self) -> Pointer;

    /// Load the descriptor table with the appropriate load instruction
    unsafe fn load(&self);
}
<commit_msg>[x86_all] Generalized `DTable::get_ptr()` to trait<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86` and `x86_64` descriptor tables (IDT, GDT, or LDT)

use memory::PAddr;
use core::mem::size_of;

/// A pointer to a descriptor table.
/// This is a format suitable
#[repr(C, packed)]
pub struct Pointer { /// the length of the descriptor table
                     pub limit: u16
                   , /// pointer to the region in memory
                     /// containing the descriptor table.
                     pub base: PAddr
                   }

/// A descriptor table (IDT or GDT)
pub trait DTable: Sized {
    /// Get the IDT pointer struct to pass to `lidt` or `lgdt`
    fn get_ptr(&self) -> Pointer {
        Pointer {
            limit: size_of::<Self>() as u16
          , base: PAddr::from(self as *const _)
        }
    }

    /// Load the descriptor table with the appropriate load instruction
    unsafe fn load(&self);
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::iter::FromIterator;

/// A very simple BitVector type.
#[derive(Clone, Debug, PartialEq)]
pub struct BitVector {
    data: Vec<u64>,
}

impl BitVector {
    #[inline]
    pub fn new(num_bits: usize) -> BitVector {
        let num_words = u64s(num_bits);
        BitVector { data: vec![0; num_words] }
    }

    #[inline]
    pub fn clear(&mut self) {
        for p in &mut self.data {
            *p = 0;
        }
    }

    pub fn count(&self) -> usize {
        self.data.iter().map(|e| e.count_ones() as usize).sum()
    }

    #[inline]
    pub fn contains(&self, bit: usize) -> bool {
        let (word, mask) = word_mask(bit);
        (self.data[word] & mask) != 0
    }

    /// Returns true if the bit has changed.
    #[inline]
    pub fn insert(&mut self, bit: usize) -> bool {
        let (word, mask) = word_mask(bit);
        let data = &mut self.data[word];
        let value = *data;
        let new_value = value | mask;
        *data = new_value;
        new_value != value
    }

    /// Returns true if the bit has changed.
    #[inline]
    pub fn remove(&mut self, bit: usize) -> bool {
        let (word, mask) = word_mask(bit);
        let data = &mut self.data[word];
        let value = *data;
        let new_value = value & !mask;
        *data = new_value;
        new_value != value
    }

    #[inline]
    pub fn insert_all(&mut self, all: &BitVector) -> bool {
        assert!(self.data.len() == all.data.len());
        let mut changed = false;
        for (i, j) in self.data.iter_mut().zip(&all.data) {
            let value = *i;
            *i = value | *j;
            if value != *i {
                changed = true;
            }
        }
        changed
    }

    #[inline]
    pub fn grow(&mut self, num_bits: usize) {
        let num_words = u64s(num_bits);
        if self.data.len() < num_words {
            self.data.resize(num_words, 0)
        }
    }

    /// Iterates over indexes of set bits in a sorted order
    #[inline]
    pub fn iter<'a>(&'a self) -> BitVectorIter<'a> {
        BitVectorIter {
            iter: self.data.iter(),
            current: 0,
            idx: 0,
        }
    }
}

pub struct BitVectorIter<'a> {
    iter: ::std::slice::Iter<'a, u64>,
    current: u64,
    idx: usize,
}

impl<'a> Iterator for BitVectorIter<'a> {
    type Item = usize;
    fn next(&mut self) -> Option<usize> {
        while self.current == 0 {
            self.current = if let Some(&i) = self.iter.next() {
                if i == 0 {
                    self.idx += 64;
                    continue;
                } else {
                    self.idx = u64s(self.idx) * 64;
                    i
                }
            } else {
                return None;
            }
        }
        let offset = self.current.trailing_zeros() as usize;
        self.current >>= offset;
        self.current >>= 1; // shift otherwise overflows for 0b1000_0000_…_0000
        self.idx += offset + 1;
        return Some(self.idx - 1);
    }
}

impl FromIterator<bool> for BitVector {
    fn from_iter<I>(iter: I) -> BitVector where I: IntoIterator<Item=bool> {
        let iter = iter.into_iter();
        let (len, _) = iter.size_hint();
        // Make the minimum length for the bitvector 64 bits since that's
        // the smallest non-zero size anyway.
        let len = if len < 64 { 64 } else { len };
        let mut bv = BitVector::new(len);
        for (idx, val) in iter.enumerate() {
            if idx > len {
                bv.grow(idx);
            }
            if val {
                bv.insert(idx);
            }
        }

        bv
    }
}

/// A "bit matrix" is basically a matrix of booleans represented as
/// one gigantic bitvector. In other words, it is as if you have
/// `rows` bitvectors, each of length `columns`.
#[derive(Clone, Debug)]
pub struct BitMatrix {
    columns: usize,
    vector: Vec<u64>,
}

impl BitMatrix {
    /// Create a new `rows x columns` matrix, initially empty.
    pub fn new(rows: usize, columns: usize) -> BitMatrix {
        // For every element, we need one bit for every other
        // element. Round up to an even number of u64s.
        let u64s_per_row = u64s(columns);
        BitMatrix {
            columns,
            vector: vec![0; rows * u64s_per_row],
        }
    }

    /// The range of bits for a given row.
    fn range(&self, row: usize) -> (usize, usize) {
        let u64s_per_row = u64s(self.columns);
        let start = row * u64s_per_row;
        (start, start + u64s_per_row)
    }

    /// Sets the cell at `(row, column)` to true. Put another way, add
    /// `column` to the bitset for `row`.
    ///
    /// Returns true if this changed the matrix, and false otherwies.
    pub fn add(&mut self, row: usize, column: usize) -> bool {
        let (start, _) = self.range(row);
        let (word, mask) = word_mask(column);
        let vector = &mut self.vector[..];
        let v1 = vector[start + word];
        let v2 = v1 | mask;
        vector[start + word] = v2;
        v1 != v2
    }

    /// Do the bits from `row` contain `column`? Put another way, is
    /// the matrix cell at `(row, column)` true?  Put yet another way,
    /// if the matrix represents (transitive) reachability, can
    /// `row` reach `column`?
    pub fn contains(&self, row: usize, column: usize) -> bool {
        let (start, _) = self.range(row);
        let (word, mask) = word_mask(column);
        (self.vector[start + word] & mask) != 0
    }

    /// Returns those indices that are true in rows `a` and `b`.  This
    /// is an O(n) operation where `n` is the number of elements
    /// (somewhat independent from the actual size of the
    /// intersection, in particular).
    pub fn intersection(&self, a: usize, b: usize) -> Vec<usize> {
        let (a_start, a_end) = self.range(a);
        let (b_start, b_end) = self.range(b);
        let mut result = Vec::with_capacity(self.columns);
        for (base, (i, j)) in (a_start..a_end).zip(b_start..b_end).enumerate() {
            let mut v = self.vector[i] & self.vector[j];
            for bit in 0..64 {
                if v == 0 {
                    break;
                }
                if v & 0x1 != 0 {
                    result.push(base * 64 + bit);
                }
                v >>= 1;
            }
        }
        result
    }

    /// Add the bits from row `read` to the bits from row `write`,
    /// return true if anything changed.
    ///
    /// This is used when computing transitive reachability because if
    /// you have an edge `write -> read`, because in that case
    /// `write` can reach everything that `read` can (and
    /// potentially more).
    pub fn merge(&mut self, read: usize, write: usize) -> bool {
        let (read_start, read_end) = self.range(read);
        let (write_start, write_end) = self.range(write);
        let vector = &mut self.vector[..];
        let mut changed = false;
        for (read_index, write_index) in (read_start..read_end).zip(write_start..write_end) {
            let v1 = vector[write_index];
            let v2 = v1 | vector[read_index];
            vector[write_index] = v2;
            changed = changed | (v1 != v2);
        }
        changed
    }

    /// Iterates through all the columns set to true in a given row of
    /// the matrix.
    pub fn iter<'a>(&'a self, row: usize) -> BitVectorIter<'a> {
        let (start, end) = self.range(row);
        BitVectorIter {
            iter: self.vector[start..end].iter(),
            current: 0,
            idx: 0,
        }
    }
}

#[inline]
fn u64s(elements: usize) -> usize {
    (elements + 63) / 64
}

#[inline]
fn word_mask(index: usize) -> (usize, u64) {
    let word = index / 64;
    let mask = 1 << (index % 64);
    (word, mask)
}

#[test]
fn bitvec_iter_works() {
    let mut bitvec = BitVector::new(100);
    bitvec.insert(1);
    bitvec.insert(10);
    bitvec.insert(19);
    bitvec.insert(62);
    bitvec.insert(63);
    bitvec.insert(64);
    bitvec.insert(65);
    bitvec.insert(66);
    bitvec.insert(99);
    assert_eq!(bitvec.iter().collect::<Vec<_>>(),
               [1, 10, 19, 62, 63, 64, 65, 66, 99]);
}

#[test]
fn bitvec_iter_works_2() {
    let mut bitvec = BitVector::new(319);
    bitvec.insert(0);
    bitvec.insert(127);
    bitvec.insert(191);
    bitvec.insert(255);
    bitvec.insert(319);
    assert_eq!(bitvec.iter().collect::<Vec<_>>(), [0, 127, 191, 255, 319]);
}

#[test]
fn union_two_vecs() {
    let mut vec1 = BitVector::new(65);
    let mut vec2 = BitVector::new(65);
    assert!(vec1.insert(3));
    assert!(!vec1.insert(3));
    assert!(vec2.insert(5));
    assert!(vec2.insert(64));
    assert!(vec1.insert_all(&vec2));
    assert!(!vec1.insert_all(&vec2));
    assert!(vec1.contains(3));
    assert!(!vec1.contains(4));
    assert!(vec1.contains(5));
    assert!(!vec1.contains(63));
    assert!(vec1.contains(64));
}

#[test]
fn grow() {
    let mut vec1 = BitVector::new(65);
    for index in 0 .. 65 {
        assert!(vec1.insert(index));
        assert!(!vec1.insert(index));
    }
    vec1.grow(128);

    // Check if the bits set before growing are still set
    for index in 0 .. 65 {
        assert!(vec1.contains(index));
    }

    // Check if the new bits are all un-set
    for index in 65 .. 128 {
        assert!(!vec1.contains(index));
    }

    // Check that we can set all new bits without running out of bounds
    for index in 65 .. 128 {
        assert!(vec1.insert(index));
        assert!(!vec1.insert(index));
    }
}

#[test]
fn matrix_intersection() {
    let mut vec1 = BitMatrix::new(200, 200);

    // (*) Elements reachable from both 2 and 65.

    vec1.add(2, 3);
    vec1.add(2, 6);
    vec1.add(2, 10); // (*)
    vec1.add(2, 64); // (*)
    vec1.add(2, 65);
    vec1.add(2, 130);
    vec1.add(2, 160); // (*)

    vec1.add(64, 133);

    vec1.add(65, 2);
    vec1.add(65, 8);
    vec1.add(65, 10); // (*)
    vec1.add(65, 64); // (*)
    vec1.add(65, 68);
    vec1.add(65, 133);
    vec1.add(65, 160); // (*)

    let intersection = vec1.intersection(2, 64);
    assert!(intersection.is_empty());

    let intersection = vec1.intersection(2, 65);
    assert_eq!(intersection, &[10, 64, 160]);
}

#[test]
fn matrix_iter() {
    let mut matrix = BitMatrix::new(64, 100);
    matrix.add(3, 22);
    matrix.add(3, 75);
    matrix.add(2, 99);
    matrix.add(4, 0);
    matrix.merge(3, 5);

    let expected = [99];
    let mut iter = expected.iter();
    for i in matrix.iter(2) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(3) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [0];
    let mut iter = expected.iter();
    for i in matrix.iter(4) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(5) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());
}
<commit_msg>Move word type and word size usage to constants & make it of 128 bits<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::iter::FromIterator;

type Word = u128;
const WORD_BITS: usize = 128;

/// A very simple BitVector type.
#[derive(Clone, Debug, PartialEq)]
pub struct BitVector {
    data: Vec<Word>,
}

impl BitVector {
    #[inline]
    pub fn new(num_bits: usize) -> BitVector {
        let num_words = words(num_bits);
        BitVector { data: vec![0; num_words] }
    }

    #[inline]
    pub fn clear(&mut self) {
        for p in &mut self.data {
            *p = 0;
        }
    }

    pub fn count(&self) -> usize {
        self.data.iter().map(|e| e.count_ones() as usize).sum()
    }

    #[inline]
    pub fn contains(&self, bit: usize) -> bool {
        let (word, mask) = word_mask(bit);
        (self.data[word] & mask) != 0
    }

    /// Returns true if the bit has changed.
    #[inline]
    pub fn insert(&mut self, bit: usize) -> bool {
        let (word, mask) = word_mask(bit);
        let data = &mut self.data[word];
        let value = *data;
        let new_value = value | mask;
        *data = new_value;
        new_value != value
    }

    /// Returns true if the bit has changed.
    #[inline]
    pub fn remove(&mut self, bit: usize) -> bool {
        let (word, mask) = word_mask(bit);
        let data = &mut self.data[word];
        let value = *data;
        let new_value = value & !mask;
        *data = new_value;
        new_value != value
    }

    #[inline]
    pub fn insert_all(&mut self, all: &BitVector) -> bool {
        assert!(self.data.len() == all.data.len());
        let mut changed = false;
        for (i, j) in self.data.iter_mut().zip(&all.data) {
            let value = *i;
            *i = value | *j;
            if value != *i {
                changed = true;
            }
        }
        changed
    }

    #[inline]
    pub fn grow(&mut self, num_bits: usize) {
        let num_words = words(num_bits);
        if self.data.len() < num_words {
            self.data.resize(num_words, 0)
        }
    }

    /// Iterates over indexes of set bits in a sorted order
    #[inline]
    pub fn iter<'a>(&'a self) -> BitVectorIter<'a> {
        BitVectorIter {
            iter: self.data.iter(),
            current: 0,
            idx: 0,
        }
    }
}

pub struct BitVectorIter<'a> {
    iter: ::std::slice::Iter<'a, Word>,
    current: Word,
    idx: usize,
}

impl<'a> Iterator for BitVectorIter<'a> {
    type Item = usize;
    fn next(&mut self) -> Option<usize> {
        while self.current == 0 {
            self.current = if let Some(&i) = self.iter.next() {
                if i == 0 {
                    self.idx += WORD_BITS;
                    continue;
                } else {
                    self.idx = words(self.idx) * WORD_BITS;
                    i
                }
            } else {
                return None;
            }
        }
        let offset = self.current.trailing_zeros() as usize;
        self.current >>= offset;
        self.current >>= 1; // shift otherwise overflows for 0b1000_0000_…_0000
        self.idx += offset + 1;
        return Some(self.idx - 1);
    }
}

impl FromIterator<bool> for BitVector {
    fn from_iter<I>(iter: I) -> BitVector where I: IntoIterator<Item=bool> {
        let iter = iter.into_iter();
        let (len, _) = iter.size_hint();
        // Make the minimum length for the bitvector WORD_BITS bits since that's
        // the smallest non-zero size anyway.
        let len = if len < WORD_BITS { WORD_BITS } else { len };
        let mut bv = BitVector::new(len);
        for (idx, val) in iter.enumerate() {
            if idx > len {
                bv.grow(idx);
            }
            if val {
                bv.insert(idx);
            }
        }

        bv
    }
}

/// A "bit matrix" is basically a matrix of booleans represented as
/// one gigantic bitvector. In other words, it is as if you have
/// `rows` bitvectors, each of length `columns`.
#[derive(Clone, Debug)]
pub struct BitMatrix {
    columns: usize,
    vector: Vec<Word>,
}

impl BitMatrix {
    /// Create a new `rows x columns` matrix, initially empty.
    pub fn new(rows: usize, columns: usize) -> BitMatrix {
        // For every element, we need one bit for every other
        // element. Round up to an even number of words.
        let words_per_row = words(columns);
        BitMatrix {
            columns,
            vector: vec![0; rows * words_per_row],
        }
    }

    /// The range of bits for a given row.
    fn range(&self, row: usize) -> (usize, usize) {
        let words_per_row = words(self.columns);
        let start = row * words_per_row;
        (start, start + words_per_row)
    }

    /// Sets the cell at `(row, column)` to true. Put another way, add
    /// `column` to the bitset for `row`.
    ///
    /// Returns true if this changed the matrix, and false otherwies.
    pub fn add(&mut self, row: usize, column: usize) -> bool {
        let (start, _) = self.range(row);
        let (word, mask) = word_mask(column);
        let vector = &mut self.vector[..];
        let v1 = vector[start + word];
        let v2 = v1 | mask;
        vector[start + word] = v2;
        v1 != v2
    }

    /// Do the bits from `row` contain `column`? Put another way, is
    /// the matrix cell at `(row, column)` true?  Put yet another way,
    /// if the matrix represents (transitive) reachability, can
    /// `row` reach `column`?
    pub fn contains(&self, row: usize, column: usize) -> bool {
        let (start, _) = self.range(row);
        let (word, mask) = word_mask(column);
        (self.vector[start + word] & mask) != 0
    }

    /// Returns those indices that are true in rows `a` and `b`.  This
    /// is an O(n) operation where `n` is the number of elements
    /// (somewhat independent from the actual size of the
    /// intersection, in particular).
    pub fn intersection(&self, a: usize, b: usize) -> Vec<usize> {
        let (a_start, a_end) = self.range(a);
        let (b_start, b_end) = self.range(b);
        let mut result = Vec::with_capacity(self.columns);
        for (base, (i, j)) in (a_start..a_end).zip(b_start..b_end).enumerate() {
            let mut v = self.vector[i] & self.vector[j];
            for bit in 0..WORD_BITS {
                if v == 0 {
                    break;
                }
                if v & 0x1 != 0 {
                    result.push(base * WORD_BITS + bit);
                }
                v >>= 1;
            }
        }
        result
    }

    /// Add the bits from row `read` to the bits from row `write`,
    /// return true if anything changed.
    ///
    /// This is used when computing transitive reachability because if
    /// you have an edge `write -> read`, because in that case
    /// `write` can reach everything that `read` can (and
    /// potentially more).
    pub fn merge(&mut self, read: usize, write: usize) -> bool {
        let (read_start, read_end) = self.range(read);
        let (write_start, write_end) = self.range(write);
        let vector = &mut self.vector[..];
        let mut changed = false;
        for (read_index, write_index) in (read_start..read_end).zip(write_start..write_end) {
            let v1 = vector[write_index];
            let v2 = v1 | vector[read_index];
            vector[write_index] = v2;
            changed = changed | (v1 != v2);
        }
        changed
    }

    /// Iterates through all the columns set to true in a given row of
    /// the matrix.
    pub fn iter<'a>(&'a self, row: usize) -> BitVectorIter<'a> {
        let (start, end) = self.range(row);
        BitVectorIter {
            iter: self.vector[start..end].iter(),
            current: 0,
            idx: 0,
        }
    }
}

#[inline]
fn words(elements: usize) -> usize {
    (elements + WORD_BITS - 1) / WORD_BITS
}

#[inline]
fn word_mask(index: usize) -> (usize, Word) {
    let word = index / WORD_BITS;
    let mask = 1 << (index % WORD_BITS);
    (word, mask)
}

#[test]
fn bitvec_iter_works() {
    let mut bitvec = BitVector::new(100);
    bitvec.insert(1);
    bitvec.insert(10);
    bitvec.insert(19);
    bitvec.insert(62);
    bitvec.insert(63);
    bitvec.insert(64);
    bitvec.insert(65);
    bitvec.insert(66);
    bitvec.insert(99);
    assert_eq!(bitvec.iter().collect::<Vec<_>>(),
               [1, 10, 19, 62, 63, 64, 65, 66, 99]);
}

#[test]
fn bitvec_iter_works_2() {
    let mut bitvec = BitVector::new(319);
    bitvec.insert(0);
    bitvec.insert(127);
    bitvec.insert(191);
    bitvec.insert(255);
    bitvec.insert(319);
    assert_eq!(bitvec.iter().collect::<Vec<_>>(), [0, 127, 191, 255, 319]);
}

#[test]
fn union_two_vecs() {
    let mut vec1 = BitVector::new(65);
    let mut vec2 = BitVector::new(65);
    assert!(vec1.insert(3));
    assert!(!vec1.insert(3));
    assert!(vec2.insert(5));
    assert!(vec2.insert(64));
    assert!(vec1.insert_all(&vec2));
    assert!(!vec1.insert_all(&vec2));
    assert!(vec1.contains(3));
    assert!(!vec1.contains(4));
    assert!(vec1.contains(5));
    assert!(!vec1.contains(63));
    assert!(vec1.contains(64));
}

#[test]
fn grow() {
    let mut vec1 = BitVector::new(65);
    for index in 0 .. 65 {
        assert!(vec1.insert(index));
        assert!(!vec1.insert(index));
    }
    vec1.grow(128);

    // Check if the bits set before growing are still set
    for index in 0 .. 65 {
        assert!(vec1.contains(index));
    }

    // Check if the new bits are all un-set
    for index in 65 .. 128 {
        assert!(!vec1.contains(index));
    }

    // Check that we can set all new bits without running out of bounds
    for index in 65 .. 128 {
        assert!(vec1.insert(index));
        assert!(!vec1.insert(index));
    }
}

#[test]
fn matrix_intersection() {
    let mut vec1 = BitMatrix::new(200, 200);

    // (*) Elements reachable from both 2 and 65.

    vec1.add(2, 3);
    vec1.add(2, 6);
    vec1.add(2, 10); // (*)
    vec1.add(2, 64); // (*)
    vec1.add(2, 65);
    vec1.add(2, 130);
    vec1.add(2, 160); // (*)

    vec1.add(64, 133);

    vec1.add(65, 2);
    vec1.add(65, 8);
    vec1.add(65, 10); // (*)
    vec1.add(65, 64); // (*)
    vec1.add(65, 68);
    vec1.add(65, 133);
    vec1.add(65, 160); // (*)

    let intersection = vec1.intersection(2, 64);
    assert!(intersection.is_empty());

    let intersection = vec1.intersection(2, 65);
    assert_eq!(intersection, &[10, 64, 160]);
}

#[test]
fn matrix_iter() {
    let mut matrix = BitMatrix::new(64, 100);
    matrix.add(3, 22);
    matrix.add(3, 75);
    matrix.add(2, 99);
    matrix.add(4, 0);
    matrix.merge(3, 5);

    let expected = [99];
    let mut iter = expected.iter();
    for i in matrix.iter(2) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(3) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [0];
    let mut iter = expected.iter();
    for i in matrix.iter(4) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(5) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());
}
<|endoftext|>


<commit_before>use addressable::Addressable;
use addressable::AddressableUtil;

// General information on 65xx: http://en.wikipedia.org/wiki/MOS_Technology_6510
// Web simulator and much info: http://e-tradition.net/bytes/6502/
// Good reference and overview: http://www.obelisk.demon.co.uk/index.html

static NMI_VECTOR: u16 = 0xfffa;
static RESET_VECTOR: u16 = 0xfffc;
static IRQ_VECTOR: u16 = 0xfffe;

struct Registers {
	pc: u16,				// program counter
	ac: u8,					// accumulator
	x: u8,					// x register
	y: u8,					// y register
	sr: u8,					// status register (NV-BDIZC: Negative, oVerflow, Break, Decimal, Interrupt, Zero, Carry)
	sp: u8,					// stack pointer
}

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn addr<M: Addressable<u16>> (&self, cpu: &Mos6502<M>) -> u16 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos6510: Immediade operand is never targetted to an address"),
			Accumulator							=> fail!("mos6510: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.reg.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.reg.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.reg.y as u16,
			Indirect(addr)						=> cpu.mem.get_le(addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.reg.x) as u16,					// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.reg.y) as u16,					// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> cpu.mem.get_le((addr + cpu.reg.x) as u16),	// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = cpu.mem.get_le(addr as u16); iaddr + cpu.reg.y as u16 },
		}
	}

	fn get<M: Addressable<u16>> (&self, cpu: &Mos6502<M>) -> u8 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.reg.ac,
			Relative(_target)					=> fail!("mos6510: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.get(addr) },
		}
	}

	fn set<M: Addressable<u16>> (&self, cpu: &mut Mos6502<M>, val: u8) {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos6510: Immediate operand never sets a value"),
			Accumulator							=> cpu.reg.ac = val,
			Relative(_target)					=> fail!("mos6510: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.set(addr, val); },
		}
	}
}

pub struct Mos6502<M> {
	priv reg: Registers,					// internal CPU registers
	priv mem: M,							// memory as accessible to the CPU
}

impl<M: Addressable<u16>> Mos6502<M> {
	pub fn new (mem: M) -> Mos6502<M> {
		Mos6502 {
			reg: Registers { pc: 0, ac: 0, x: 0, y: 0, sr: 0, sp: 0 },
			mem: mem,
		}
	}

	fn get_opcode (&self) -> u8 {
		self.mem.get(self.reg.pc)
	}

	fn get_argument<T: Int> (&self) -> T {
		self.mem.get_le(self.reg.pc + 1)
	}

	pub fn step (&mut self) -> uint {
		// TODO
		0
	}
}

pub struct Mos6510<M> {
	priv cpu: Mos6502<M>,					// Core CPU
	priv port_ddr: u8,						// CPU port data direction register
	priv port_dat: u8,						// CPU port data register
}

impl<M: Addressable<u16>> Mos6510<M> {
	pub fn new (mem: M) -> Mos6510<M> {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		Mos6510 { cpu: Mos6502::new(mem), port_ddr: 0, port_dat: 0 }
	}

	pub fn step (&mut self) -> uint {
		self.cpu.step()
	}
}
<commit_msg>Moved registers back into Mos6502<commit_after>use addressable::Addressable;
use addressable::AddressableUtil;

// General information on 65xx: http://en.wikipedia.org/wiki/MOS_Technology_6510
// Web simulator and much info: http://e-tradition.net/bytes/6502/
// Good reference and overview: http://www.obelisk.demon.co.uk/index.html

static NMI_VECTOR: u16 = 0xfffa;
static RESET_VECTOR: u16 = 0xfffc;
static IRQ_VECTOR: u16 = 0xfffe;

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn addr<M: Addressable<u16>> (&self, cpu: &Mos6502<M>) -> u16 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos6510: Immediade operand is never targetted to an address"),
			Accumulator							=> fail!("mos6510: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.y as u16,
			Indirect(addr)						=> cpu.mem.get_le(addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.x) as u16,						// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.y) as u16,						// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> cpu.mem.get_le((addr + cpu.x) as u16),		// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = cpu.mem.get_le(addr as u16); iaddr + cpu.y as u16 },
		}
	}

	fn get<M: Addressable<u16>> (&self, cpu: &Mos6502<M>) -> u8 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.ac,
			Relative(_target)					=> fail!("mos6510: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.get(addr) },
		}
	}

	fn set<M: Addressable<u16>> (&self, cpu: &mut Mos6502<M>, val: u8) {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos6510: Immediate operand never sets a value"),
			Accumulator							=> cpu.ac = val,
			Relative(_target)					=> fail!("mos6510: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.set(addr, val); },
		}
	}
}

pub struct Mos6502 {
	priv pc: u16,						// program counter
	priv ac: u8,						// accumulator
	priv x: u8,							// x register
	priv y: u8,							// y register
	priv sr: u8,						// status register (NV-BDIZC: Negative, oVerflow, Break, Decimal, Interrupt, Zero, Carry)
	priv sp: u8,						// stack pointer
}

impl<M: Addressable<u16>> Mos6502<M> {
	pub fn new (mem: M) -> Mos6502<M> {
		Mos6502 {
			pc: 0,
			ac: 0,
			x: 0,
			y: 0,
			sr: 0,
			sp: 0,
		}
	}

	fn get_opcode (&self) -> u8 {
		self.mem.get(self.reg.pc)
	}

	fn get_argument<T: Int> (&self) -> T {
		self.mem.get_le(self.reg.pc + 1)
	}

	pub fn step (&mut self) -> uint {
		// TODO
		0
	}
}

pub struct Mos6510<M> {
	priv cpu: Mos6502<M>,					// Core CPU
	priv port_ddr: u8,						// CPU port data direction register
	priv port_dat: u8,						// CPU port data register
}

impl<M: Addressable<u16>> Mos6510<M> {
	pub fn new (mem: M) -> Mos6510<M> {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		Mos6510 { cpu: Mos6502::new(mem), port_ddr: 0, port_dat: 0 }
	}

	pub fn step (&mut self) -> uint {
		self.cpu.step()
	}
}
<|endoftext|>


<commit_before>use engine::Game;
use gui::{Console, Key, Menu};
use gui::screens::{self, Screen, ScreenChange};
use util::units::Point;

use std::fmt;

#[allow(missing_copy_implementations)]
pub struct MenuScreen {
    menu: Menu<MenuItem>,
}

static LOGO: &'static str =
       "                                 \n\
        ##### ##### #     #  #   # ##### \n\
        #   # #     #     #  #  #  #     \n\
        #   # #     #     #  # #   #     \n\
        ##### ##### #     #  ##    ###   \n\
        # #       # #     #  # #   #     \n\
        #  #      # #     #  #  #  #     \n\
        #   # ##### ##### #  #   # ##### \n";

enum MenuItem {
    StartGame,
    Exit,
}

impl fmt::Display for MenuItem {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MenuItem::StartGame => write!(fmt, "Start Game"),
            MenuItem::Exit => write!(fmt, "Exit Game"),
        }
    }
}

impl MenuScreen {
    pub fn new() -> Box<Screen + 'static> {
        Box::new(
            MenuScreen {
                menu: Menu::new(vec![MenuItem::StartGame, MenuItem::Exit]),
            }
        )
    }
}

impl Screen for MenuScreen {

    #[allow(unused)]
    fn input(&mut self, game: &mut Game, console: &mut Console) -> Option<ScreenChange> {
        if let Some(key) = console.check_for_keypress() {
            match key {
                Key::Up => {
                    self.menu.prev();
                }
                Key::Down => {
                    self.menu.next();
                }
                Key::Enter => {
                    match *self.menu.selected() {
                        MenuItem::StartGame => return Some(ScreenChange::AddScreen(screens::GameScreen::new())),
                        MenuItem::Exit => return Some(ScreenChange::ExitGame),
                    }
                },
                _ => {}
            }
        }

        None
    }

    #[allow(unused)]
    fn update(&mut self, game: &mut Game, console: &mut Console) -> Option<ScreenChange> {
        None
    }

    #[allow(unused)]
    fn render(&mut self, game: &mut Game, console: &mut Console) {
        let logo_loc = Point::new(10, 10);
        console.print_plain(logo_loc, LOGO);

        let menu_loc = Point::new(20, 20);

        for (i, menu_item) in self.menu.enum_items() {
            console.print_plain(menu_loc.right(2).down(i as i32), format!("{}", menu_item).as_slice());
            if self.menu.is_selected(i) {
                console.put_plain(menu_loc.down(i as i32), '>');
            }
        }
    }
}
<commit_msg>Centered items on the menu screen<commit_after>use engine::Game;
use gui::{Console, Key, Menu};
use gui::screens::{self, Screen, ScreenChange};
use util::units::Point;

use std::fmt;

#[allow(missing_copy_implementations)]
pub struct MenuScreen {
    menu: Menu<MenuItem>,
}

static LOGO: &'static str =
       "                                 \n\
        ##### ##### #     #  #   # ##### \n\
        #   # #     #     #  #  #  #     \n\
        #   # #     #     #  # #   #     \n\
        ##### ##### #     #  ##    ###   \n\
        # #       # #     #  # #   #     \n\
        #  #      # #     #  #  #  #     \n\
        #   # ##### ##### #  #   # ##### \n";

enum MenuItem {
    StartGame,
    Exit,
}

impl fmt::Display for MenuItem {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MenuItem::StartGame => write!(fmt, "Start Game"),
            MenuItem::Exit => write!(fmt, "Exit Game"),
        }
    }
}

impl MenuScreen {
    pub fn new() -> Box<Screen + 'static> {
        Box::new(
            MenuScreen {
                menu: Menu::new(vec![MenuItem::StartGame, MenuItem::Exit]),
            }
        )
    }
}

impl Screen for MenuScreen {

    #[allow(unused)]
    fn input(&mut self, game: &mut Game, console: &mut Console) -> Option<ScreenChange> {
        if let Some(key) = console.check_for_keypress() {
            match key {
                Key::Up => {
                    self.menu.prev();
                }
                Key::Down => {
                    self.menu.next();
                }
                Key::Enter => {
                    match *self.menu.selected() {
                        MenuItem::StartGame => return Some(ScreenChange::AddScreen(screens::GameScreen::new())),
                        MenuItem::Exit => return Some(ScreenChange::ExitGame),
                    }
                },
                _ => {}
            }
        }

        None
    }

    #[allow(unused)]
    fn update(&mut self, game: &mut Game, console: &mut Console) -> Option<ScreenChange> {
        None
    }

    #[allow(unused)]
    fn render(&mut self, game: &mut Game, console: &mut Console) {
        let logo_loc = Point::new(24, 20);
        console.print_plain(logo_loc, LOGO);

        let menu_loc = Point::new(33, 30);

        for (i, menu_item) in self.menu.enum_items() {
            console.print_plain(menu_loc.right(2).down(i as i32), format!("{}", menu_item).as_slice());
            if self.menu.is_selected(i) {
                console.put_plain(menu_loc.down(i as i32), '>');
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Timers for non-linux/non-windows OSes
//!
//! This module implements timers with a worker thread, select(), and a lot of
//! witchcraft that turns out to be horribly inaccurate timers. The unfortunate
//! part is that I'm at a loss of what else to do one these OSes. This is also
//! why linux has a specialized timerfd implementation and windows has its own
//! implementation (they're more accurate than this one).
//!
//! The basic idea is that there is a worker thread that's communicated to via a
//! channel and a pipe, the pipe is used by the worker thread in a select()
//! syscall with a timeout. The timeout is the "next timer timeout" while the
//! channel is used to send data over to the worker thread.
//!
//! Whenever the call to select() times out, then a channel receives a message.
//! Whenever the call returns that the file descriptor has information, then the
//! channel from timers is drained, enqueueing all incoming requests.
//!
//! The actual implementation of the helper thread is a sorted array of
//! timers in terms of target firing date. The target is the absolute time at
//! which the timer should fire. Timers are then re-enqueued after a firing if
//! the repeat boolean is set.
//!
//! Naturally, all this logic of adding times and keeping track of
//! relative/absolute time is a little lossy and not quite exact. I've done the
//! best I could to reduce the amount of calls to 'now()', but there's likely
//! still inaccuracies trickling in here and there.
//!
//! One of the tricky parts of this implementation is that whenever a timer is
//! acted upon, it must cancel whatever the previous action was (if one is
//! active) in order to act like the other implementations of this timer. In
//! order to do this, the timer's inner pointer is transferred to the worker
//! thread. Whenever the timer is modified, it first takes ownership back from
//! the worker thread in order to modify the same data structure. This has the
//! side effect of "cancelling" the previous requests while allowing a
//! re-enqueueing later on.
//!
//! Note that all time units in this file are in *milliseconds*.

use libc;
use std::mem;
use std::os;
use std::ptr;
use std::rt::rtio;
use std::sync::atomics;
use std::comm;

use io::IoResult;
use io::c;
use io::file::FileDesc;
use io::helper_thread::Helper;

helper_init!(static mut HELPER: Helper<Req>)

pub struct Timer {
    id: uint,
    inner: Option<Box<Inner>>,
}

struct Inner {
    tx: Option<Sender<()>>,
    interval: u64,
    repeat: bool,
    target: u64,
    id: uint,
}

#[allow(visible_private_types)]
pub enum Req {
    // Add a new timer to the helper thread.
    NewTimer(Box<Inner>),

    // Remove a timer based on its id and then send it back on the channel
    // provided
    RemoveTimer(uint, Sender<Box<Inner>>),
}

// returns the current time (in milliseconds)
pub fn now() -> u64 {
    unsafe {
        let mut now: libc::timeval = mem::zeroed();
        assert_eq!(c::gettimeofday(&mut now, ptr::null()), 0);
        return (now.tv_sec as u64) * 1000 + (now.tv_usec as u64) / 1000;
    }
}

fn helper(input: libc::c_int, messages: Receiver<Req>, _: ()) {
    let mut set: c::fd_set = unsafe { mem::zeroed() };

    let mut fd = FileDesc::new(input, true);
    let mut timeout: libc::timeval = unsafe { mem::zeroed() };

    // active timers are those which are able to be selected upon (and it's a
    // sorted list, and dead timers are those which have expired, but ownership
    // hasn't yet been transferred back to the timer itself.
    let mut active: Vec<Box<Inner>> = vec![];
    let mut dead = vec![];

    // inserts a timer into an array of timers (sorted by firing time)
    fn insert(t: Box<Inner>, active: &mut Vec<Box<Inner>>) {
        match active.iter().position(|tm| tm.target > t.target) {
            Some(pos) => { active.insert(pos, t); }
            None => { active.push(t); }
        }
    }

    // signals the first requests in the queue, possible re-enqueueing it.
    fn signal(active: &mut Vec<Box<Inner>>,
              dead: &mut Vec<(uint, Box<Inner>)>) {
        let mut timer = match active.shift() {
            Some(timer) => timer, None => return
        };
        let tx = timer.tx.take_unwrap();
        if tx.send_opt(()).is_ok() && timer.repeat {
            timer.tx = Some(tx);
            timer.target += timer.interval;
            insert(timer, active);
        } else {
            drop(tx);
            dead.push((timer.id, timer));
        }
    }

    'outer: loop {
        let timeout = if active.len() == 0 {
            // Empty array? no timeout (wait forever for the next request)
            ptr::null()
        } else {
            let now = now();
            // If this request has already expired, then signal it and go
            // through another iteration
            if active.get(0).target <= now {
                signal(&mut active, &mut dead);
                continue;
            }

            // The actual timeout listed in the requests array is an
            // absolute date, so here we translate the absolute time to a
            // relative time.
            let tm = active.get(0).target - now;
            timeout.tv_sec = (tm / 1000) as libc::time_t;
            timeout.tv_usec = ((tm % 1000) * 1000) as libc::suseconds_t;
            &timeout as *libc::timeval
        };

        c::fd_set(&mut set, input);
        match unsafe {
            c::select(input + 1, &set, ptr::null(), ptr::null(), timeout)
        } {
            // timed out
            0 => signal(&mut active, &mut dead),

            // file descriptor write woke us up, we've got some new requests
            1 => {
                loop {
                    match messages.try_recv() {
                        Err(comm::Disconnected) => {
                            assert!(active.len() == 0);
                            break 'outer;
                        }

                        Ok(NewTimer(timer)) => insert(timer, &mut active),

                        Ok(RemoveTimer(id, ack)) => {
                            match dead.iter().position(|&(i, _)| id == i) {
                                Some(i) => {
                                    let (_, i) = dead.remove(i).unwrap();
                                    ack.send(i);
                                    continue
                                }
                                None => {}
                            }
                            let i = active.iter().position(|i| i.id == id);
                            let i = i.expect("no timer found");
                            let t = active.remove(i).unwrap();
                            ack.send(t);
                        }
                        Err(..) => break
                    }
                }

                // drain the file descriptor
                let mut buf = [0];
                assert_eq!(fd.inner_read(buf).unwrap(), 1);
            }

            -1 if os::errno() == libc::EINTR as int => {}
            n => fail!("helper thread failed in select() with error: {} ({})",
                       n, os::last_os_error())
        }
    }
}

impl Timer {
    pub fn new() -> IoResult<Timer> {
        unsafe { HELPER.boot(|| {}, helper); }

        static mut ID: atomics::AtomicUint = atomics::INIT_ATOMIC_UINT;
        let id = unsafe { ID.fetch_add(1, atomics::Relaxed) };
        Ok(Timer {
            id: id,
            inner: Some(box Inner {
                tx: None,
                interval: 0,
                target: 0,
                repeat: false,
                id: id,
            })
        })
    }

    pub fn sleep(ms: u64) {
        let mut to_sleep = libc::timespec {
            tv_sec: (ms / 1000) as libc::time_t,
            tv_nsec: ((ms % 1000) * 1000000) as libc::c_long,
        };
        while unsafe { libc::nanosleep(&to_sleep, &mut to_sleep) } != 0 {
            if os::errno() as int != libc::EINTR as int {
                fail!("failed to sleep, but not because of EINTR?");
            }
        }
    }

    fn inner(&mut self) -> Box<Inner> {
        match self.inner.take() {
            Some(i) => i,
            None => {
                let (tx, rx) = channel();
                unsafe { HELPER.send(RemoveTimer(self.id, tx)); }
                rx.recv()
            }
        }
    }
}

impl rtio::RtioTimer for Timer {
    fn sleep(&mut self, msecs: u64) {
        let mut inner = self.inner();
        inner.tx = None; // cancel any previous request
        self.inner = Some(inner);

        Timer::sleep(msecs);
    }

    fn oneshot(&mut self, msecs: u64) -> Receiver<()> {
        let now = now();
        let mut inner = self.inner();

        let (tx, rx) = channel();
        inner.repeat = false;
        inner.tx = Some(tx);
        inner.interval = msecs;
        inner.target = now + msecs;

        unsafe { HELPER.send(NewTimer(inner)); }
        return rx;
    }

    fn period(&mut self, msecs: u64) -> Receiver<()> {
        let now = now();
        let mut inner = self.inner();

        let (tx, rx) = channel();
        inner.repeat = true;
        inner.tx = Some(tx);
        inner.interval = msecs;
        inner.target = now + msecs;

        unsafe { HELPER.send(NewTimer(inner)); }
        return rx;
    }
}

impl Drop for Timer {
    fn drop(&mut self) {
        self.inner = Some(self.inner());
    }
}
<commit_msg>!FIXME: ARM compilation workaround<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Timers for non-linux/non-windows OSes
//!
//! This module implements timers with a worker thread, select(), and a lot of
//! witchcraft that turns out to be horribly inaccurate timers. The unfortunate
//! part is that I'm at a loss of what else to do one these OSes. This is also
//! why linux has a specialized timerfd implementation and windows has its own
//! implementation (they're more accurate than this one).
//!
//! The basic idea is that there is a worker thread that's communicated to via a
//! channel and a pipe, the pipe is used by the worker thread in a select()
//! syscall with a timeout. The timeout is the "next timer timeout" while the
//! channel is used to send data over to the worker thread.
//!
//! Whenever the call to select() times out, then a channel receives a message.
//! Whenever the call returns that the file descriptor has information, then the
//! channel from timers is drained, enqueueing all incoming requests.
//!
//! The actual implementation of the helper thread is a sorted array of
//! timers in terms of target firing date. The target is the absolute time at
//! which the timer should fire. Timers are then re-enqueued after a firing if
//! the repeat boolean is set.
//!
//! Naturally, all this logic of adding times and keeping track of
//! relative/absolute time is a little lossy and not quite exact. I've done the
//! best I could to reduce the amount of calls to 'now()', but there's likely
//! still inaccuracies trickling in here and there.
//!
//! One of the tricky parts of this implementation is that whenever a timer is
//! acted upon, it must cancel whatever the previous action was (if one is
//! active) in order to act like the other implementations of this timer. In
//! order to do this, the timer's inner pointer is transferred to the worker
//! thread. Whenever the timer is modified, it first takes ownership back from
//! the worker thread in order to modify the same data structure. This has the
//! side effect of "cancelling" the previous requests while allowing a
//! re-enqueueing later on.
//!
//! Note that all time units in this file are in *milliseconds*.

use libc;
use std::mem;
use std::os;
use std::ptr;
use std::rt::rtio;
use std::sync::atomics;
use std::comm;

use io::IoResult;
use io::c;
use io::file::FileDesc;
use io::helper_thread::Helper;

helper_init!(static mut HELPER: Helper<Req>)

pub struct Timer {
    id: uint,
    inner: Option<Box<Inner>>,
}

struct Inner {
    tx: Option<Sender<()>>,
    interval: u64,
    repeat: bool,
    target: u64,
    id: uint,
}

#[allow(visible_private_types)]
pub enum Req {
    // Add a new timer to the helper thread.
    NewTimer(Box<Inner>),

    // Remove a timer based on its id and then send it back on the channel
    // provided
    RemoveTimer(uint, Sender<Box<Inner>>),
}

// returns the current time (in milliseconds)
pub fn now() -> u64 {
    unsafe {
        let mut now: libc::timeval = mem::zeroed();
        assert_eq!(c::gettimeofday(&mut now, ptr::null()), 0);
        return (now.tv_sec as u64) * 1000 + (now.tv_usec as u64) / 1000;
    }
}

fn helper(input: libc::c_int, messages: Receiver<Req>, _: int) {
    let mut set: c::fd_set = unsafe { mem::zeroed() };

    let mut fd = FileDesc::new(input, true);
    let mut timeout: libc::timeval = unsafe { mem::zeroed() };

    // active timers are those which are able to be selected upon (and it's a
    // sorted list, and dead timers are those which have expired, but ownership
    // hasn't yet been transferred back to the timer itself.
    let mut active: Vec<Box<Inner>> = vec![];
    let mut dead = vec![];

    // inserts a timer into an array of timers (sorted by firing time)
    fn insert(t: Box<Inner>, active: &mut Vec<Box<Inner>>) {
        match active.iter().position(|tm| tm.target > t.target) {
            Some(pos) => { active.insert(pos, t); }
            None => { active.push(t); }
        }
    }

    // signals the first requests in the queue, possible re-enqueueing it.
    fn signal(active: &mut Vec<Box<Inner>>,
              dead: &mut Vec<(uint, Box<Inner>)>) {
        let mut timer = match active.shift() {
            Some(timer) => timer, None => return
        };
        let tx = timer.tx.take_unwrap();
        if tx.send_opt(()).is_ok() && timer.repeat {
            timer.tx = Some(tx);
            timer.target += timer.interval;
            insert(timer, active);
        } else {
            drop(tx);
            dead.push((timer.id, timer));
        }
    }

    'outer: loop {
        let timeout = if active.len() == 0 {
            // Empty array? no timeout (wait forever for the next request)
            ptr::null()
        } else {
            let now = now();
            // If this request has already expired, then signal it and go
            // through another iteration
            if active.get(0).target <= now {
                signal(&mut active, &mut dead);
                continue;
            }

            // The actual timeout listed in the requests array is an
            // absolute date, so here we translate the absolute time to a
            // relative time.
            let tm = active.get(0).target - now;
            timeout.tv_sec = (tm / 1000) as libc::time_t;
            timeout.tv_usec = ((tm % 1000) * 1000) as libc::suseconds_t;
            &timeout as *libc::timeval
        };

        c::fd_set(&mut set, input);
        match unsafe {
            c::select(input + 1, &set, ptr::null(), ptr::null(), timeout)
        } {
            // timed out
            0 => signal(&mut active, &mut dead),

            // file descriptor write woke us up, we've got some new requests
            1 => {
                loop {
                    match messages.try_recv() {
                        Err(comm::Disconnected) => {
                            assert!(active.len() == 0);
                            break 'outer;
                        }

                        Ok(NewTimer(timer)) => insert(timer, &mut active),

                        Ok(RemoveTimer(id, ack)) => {
                            match dead.iter().position(|&(i, _)| id == i) {
                                Some(i) => {
                                    let (_, i) = dead.remove(i).unwrap();
                                    ack.send(i);
                                    continue
                                }
                                None => {}
                            }
                            let i = active.iter().position(|i| i.id == id);
                            let i = i.expect("no timer found");
                            let t = active.remove(i).unwrap();
                            ack.send(t);
                        }
                        Err(..) => break
                    }
                }

                // drain the file descriptor
                let mut buf = [0];
                assert_eq!(fd.inner_read(buf).unwrap(), 1);
            }

            -1 if os::errno() == libc::EINTR as int => {}
            n => fail!("helper thread failed in select() with error: {} ({})",
                       n, os::last_os_error())
        }
    }
}

impl Timer {
    pub fn new() -> IoResult<Timer> {
        unsafe { HELPER.boot(|| {0}, helper); }

        static mut ID: atomics::AtomicUint = atomics::INIT_ATOMIC_UINT;
        let id = unsafe { ID.fetch_add(1, atomics::Relaxed) };
        Ok(Timer {
            id: id,
            inner: Some(box Inner {
                tx: None,
                interval: 0,
                target: 0,
                repeat: false,
                id: id,
            })
        })
    }

    pub fn sleep(ms: u64) {
        let mut to_sleep = libc::timespec {
            tv_sec: (ms / 1000) as libc::time_t,
            tv_nsec: ((ms % 1000) * 1000000) as libc::c_long,
        };
        while unsafe { libc::nanosleep(&to_sleep, &mut to_sleep) } != 0 {
            if os::errno() as int != libc::EINTR as int {
                fail!("failed to sleep, but not because of EINTR?");
            }
        }
    }

    fn inner(&mut self) -> Box<Inner> {
        match self.inner.take() {
            Some(i) => i,
            None => {
                let (tx, rx) = channel();
                unsafe { HELPER.send(RemoveTimer(self.id, tx)); }
                rx.recv()
            }
        }
    }
}

impl rtio::RtioTimer for Timer {
    fn sleep(&mut self, msecs: u64) {
        let mut inner = self.inner();
        inner.tx = None; // cancel any previous request
        self.inner = Some(inner);

        Timer::sleep(msecs);
    }

    fn oneshot(&mut self, msecs: u64) -> Receiver<()> {
        let now = now();
        let mut inner = self.inner();

        let (tx, rx) = channel();
        inner.repeat = false;
        inner.tx = Some(tx);
        inner.interval = msecs;
        inner.target = now + msecs;

        unsafe { HELPER.send(NewTimer(inner)); }
        return rx;
    }

    fn period(&mut self, msecs: u64) -> Receiver<()> {
        let now = now();
        let mut inner = self.inner();

        let (tx, rx) = channel();
        inner.repeat = true;
        inner.tx = Some(tx);
        inner.interval = msecs;
        inner.target = now + msecs;

        unsafe { HELPER.send(NewTimer(inner)); }
        return rx;
    }
}

impl Drop for Timer {
    fn drop(&mut self) {
        self.inner = Some(self.inner());
    }
}
<|endoftext|>