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

use std::sync::{mpsc, Arc, RwLock};
use std::time::{Duration, Instant};
use std::thread::{self, JoinHandle};

use linked_hash_map::LinkedHashMap;
use hab_net::config::ToAddrString;
use hab_net::server::ZMQ_CONTEXT;
use protobuf::{parse_from_bytes, Message};
use protocol::jobsrv;
use zmq;

use config::Config;
use data_store::DataStore;
use error::Result;

const WORKER_MGR_ADDR: &'static str = "inproc://work-manager";
const WORKER_TIMEOUT_MS: u64 = 33_000;

pub struct WorkerMgrClient {
    socket: zmq::Socket,
}

impl WorkerMgrClient {
    pub fn connect(&mut self) -> Result<()> {
        try!(self.socket.connect(WORKER_MGR_ADDR));
        Ok(())
    }

    pub fn notify_work(&mut self) -> Result<()> {
        try!(self.socket.send(&[1], 0));
        Ok(())
    }
}

impl Default for WorkerMgrClient {
    fn default() -> WorkerMgrClient {
        let socket = (**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER).unwrap();
        socket.set_sndhwm(1).unwrap();
        socket.set_linger(0).unwrap();
        socket.set_immediate(true).unwrap();
        WorkerMgrClient { socket: socket }
    }
}

pub struct WorkerMgr {
    config: Arc<RwLock<Config>>,
    datastore: DataStore,
    hb_sock: zmq::Socket,
    rq_sock: zmq::Socket,
    work_mgr_sock: zmq::Socket,
    msg: zmq::Message,
    workers: LinkedHashMap<String, Instant>,
}

impl WorkerMgr {
    pub fn new(config: Arc<RwLock<Config>>, datastore: DataStore) -> Result<Self> {
        let hb_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::SUB));
        let rq_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::ROUTER));
        let work_mgr_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER));
        try!(rq_sock.set_router_mandatory(true));
        try!(hb_sock.set_subscribe(&[]));
        try!(work_mgr_sock.set_rcvhwm(1));
        try!(work_mgr_sock.set_linger(0));
        try!(work_mgr_sock.set_immediate(true));
        let msg = try!(zmq::Message::new());
        Ok(WorkerMgr {
            config: config,
            datastore: datastore,
            hb_sock: hb_sock,
            rq_sock: rq_sock,
            work_mgr_sock: work_mgr_sock,
            msg: msg,
            workers: LinkedHashMap::new(),
        })
    }

    pub fn start(cfg: Arc<RwLock<Config>>, ds: DataStore) -> Result<JoinHandle<()>> {
        let (tx, rx) = mpsc::sync_channel(1);
        let handle = thread::Builder::new()
            .name("worker-manager".to_string())
            .spawn(move || {
                let mut manager = Self::new(cfg, ds).unwrap();
                manager.run(tx).unwrap();
            })
            .unwrap();
        match rx.recv() {
            Ok(()) => Ok(handle),
            Err(e) => panic!("worker-manager thread startup error, err={}", e),
        }
    }

    fn run(&mut self, rz: mpsc::SyncSender<()>) -> Result<()> {
        try!(self.work_mgr_sock.bind(WORKER_MGR_ADDR));
        {
            let cfg = self.config.read().unwrap();
            println!("Listening for commands on {}",
                     cfg.worker_command_addr.to_addr_string());
            try!(self.rq_sock.bind(&cfg.worker_command_addr.to_addr_string()));
            println!("Listening for heartbeats on {}",
                     cfg.worker_heartbeat_addr.to_addr_string());
            try!(self.hb_sock.bind(&cfg.worker_heartbeat_addr.to_addr_string()));
        }
        let mut hb_sock = false;
        let mut rq_sock = false;
        let mut work_mgr_sock = false;
        rz.send(()).unwrap();
        loop {
            {
                let timeout = self.poll_timeout();
                let mut items = [self.hb_sock.as_poll_item(1),
                                 self.rq_sock.as_poll_item(1),
                                 self.work_mgr_sock.as_poll_item(1)];
                // Poll until timeout or message is received. Checking for the zmq::POLLIN flag on
                // a poll item's revents will let you know if you have received a message or not
                // on that socket.
                try!(zmq::poll(&mut items, timeout));
                if (items[0].get_revents() & zmq::POLLIN) > 0 {
                    hb_sock = true;
                }
                if (items[1].get_revents() & zmq::POLLIN) > 0 {
                    rq_sock = true;
                }
                if (items[2].get_revents() & zmq::POLLIN) > 0 {
                    work_mgr_sock = true;
                }
            }
            if hb_sock {
                try!(self.process_heartbeat());
                hb_sock = false;
            }
            self.expire_workers();
            if rq_sock {
                try!(self.process_job_status());
                rq_sock = false;
            }
            if work_mgr_sock {
                try!(self.distribute_work());
            }
        }
    }

    fn poll_timeout(&self) -> i64 {
        if let Some((_, expiry)) = self.workers.front() {
            let timeout = *expiry - Instant::now();
            (timeout.as_secs() as i64 * 1000) + (timeout.subsec_nanos() as i64 / 1000 / 1000)
        } else {
            -1
        }
    }

    fn distribute_work(&mut self) -> Result<()> {
        loop {
            // Take one job from the pending list
            let mut jobs = self.datastore.pending_jobs(1)?;
            // 0 means there are no pending jobs, so we should consume our notice that we have work
            if jobs.len() == 0 {
                try!(self.work_mgr_sock.recv(&mut self.msg, 0));
                break;
            }
            // This unwrap is fine, because we just checked our length
            let mut job = jobs.pop().unwrap();

            match self.workers.pop_front() {
                Some((worker, _)) => {
                    debug!("sending work, worker={:?}, job={:?}", worker, job);
                    if self.rq_sock.send_str(&worker, zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.set_job_state(&job)?;
                        continue;
                    }
                    if self.rq_sock.send(&[], zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.set_job_state(&job)?;
                        continue;
                    }
                    if self.rq_sock.send(&job.write_to_bytes().unwrap(), 0).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.set_job_state(&job)?;
                        continue;
                    }
                }
                None => break,
            }
        }
        Ok(())
    }

    fn expire_workers(&mut self) {
        let now = Instant::now();
        loop {
            if let Some((_, expiry)) = self.workers.front() {
                if expiry >= &now {
                    break;
                }
            } else {
                break;
            }
            let worker = self.workers.pop_front();
            debug!("expiring worker due to inactivity, worker={:?}", worker);
        }
    }

    fn process_heartbeat(&mut self) -> Result<()> {
        try!(self.hb_sock.recv(&mut self.msg, 0));
        let heartbeat: jobsrv::Heartbeat = try!(parse_from_bytes(&self.msg));
        debug!("heartbeat={:?}", heartbeat);
        match heartbeat.get_state() {
            jobsrv::WorkerState::Ready => {
                let now = Instant::now();
                let expiry = now + Duration::from_millis(WORKER_TIMEOUT_MS);
                self.workers.insert(heartbeat.get_endpoint().to_string(), expiry);
            }
            jobsrv::WorkerState::Busy => {
                self.workers.remove(heartbeat.get_endpoint());
            }
        }
        Ok(())
    }

    fn process_job_status(&mut self) -> Result<()> {
        // Pop message delimiter
        try!(self.rq_sock.recv(&mut self.msg, 0));
        // Pop message body
        try!(self.rq_sock.recv(&mut self.msg, 0));
        let job: jobsrv::Job = try!(parse_from_bytes(&self.msg));
        debug!("job_status={:?}", job);
        try!(self.datastore.set_job_state(&job));
        Ok(())
    }
}
<commit_msg>Fix deadlock in job server<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::{mpsc, Arc, RwLock};
use std::time::{Duration, Instant};
use std::thread::{self, JoinHandle};

use linked_hash_map::LinkedHashMap;
use hab_net::config::ToAddrString;
use hab_net::server::ZMQ_CONTEXT;
use protobuf::{parse_from_bytes, Message};
use protocol::jobsrv;
use zmq;

use config::Config;
use data_store::DataStore;
use error::Result;

const WORKER_MGR_ADDR: &'static str = "inproc://work-manager";
const WORKER_TIMEOUT_MS: u64 = 33_000;

pub struct WorkerMgrClient {
    socket: zmq::Socket,
}

impl WorkerMgrClient {
    pub fn connect(&mut self) -> Result<()> {
        try!(self.socket.connect(WORKER_MGR_ADDR));
        Ok(())
    }

    pub fn notify_work(&mut self) -> Result<()> {
        try!(self.socket.send(&[1], 0));
        Ok(())
    }
}

impl Default for WorkerMgrClient {
    fn default() -> WorkerMgrClient {
        let socket = (**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER).unwrap();
        socket.set_sndhwm(1).unwrap();
        socket.set_linger(0).unwrap();
        socket.set_immediate(true).unwrap();
        WorkerMgrClient { socket: socket }
    }
}

pub struct WorkerMgr {
    config: Arc<RwLock<Config>>,
    datastore: DataStore,
    hb_sock: zmq::Socket,
    rq_sock: zmq::Socket,
    work_mgr_sock: zmq::Socket,
    msg: zmq::Message,
    workers: LinkedHashMap<String, Instant>,
}

impl WorkerMgr {
    pub fn new(config: Arc<RwLock<Config>>, datastore: DataStore) -> Result<Self> {
        let hb_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::SUB));
        let rq_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::ROUTER));
        let work_mgr_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER));
        try!(rq_sock.set_router_mandatory(true));
        try!(hb_sock.set_subscribe(&[]));
        try!(work_mgr_sock.set_rcvhwm(1));
        try!(work_mgr_sock.set_linger(0));
        try!(work_mgr_sock.set_immediate(true));
        let msg = try!(zmq::Message::new());
        Ok(WorkerMgr {
            config: config,
            datastore: datastore,
            hb_sock: hb_sock,
            rq_sock: rq_sock,
            work_mgr_sock: work_mgr_sock,
            msg: msg,
            workers: LinkedHashMap::new(),
        })
    }

    pub fn start(cfg: Arc<RwLock<Config>>, ds: DataStore) -> Result<JoinHandle<()>> {
        let (tx, rx) = mpsc::sync_channel(1);
        let handle = thread::Builder::new()
            .name("worker-manager".to_string())
            .spawn(move || {
                let mut manager = Self::new(cfg, ds).unwrap();
                manager.run(tx).unwrap();
            })
            .unwrap();
        match rx.recv() {
            Ok(()) => Ok(handle),
            Err(e) => panic!("worker-manager thread startup error, err={}", e),
        }
    }

    fn run(&mut self, rz: mpsc::SyncSender<()>) -> Result<()> {
        try!(self.work_mgr_sock.bind(WORKER_MGR_ADDR));
        {
            let cfg = self.config.read().unwrap();
            println!("Listening for commands on {}",
                     cfg.worker_command_addr.to_addr_string());
            try!(self.rq_sock.bind(&cfg.worker_command_addr.to_addr_string()));
            println!("Listening for heartbeats on {}",
                     cfg.worker_heartbeat_addr.to_addr_string());
            try!(self.hb_sock.bind(&cfg.worker_heartbeat_addr.to_addr_string()));
        }
        let mut hb_sock = false;
        let mut rq_sock = false;
        let mut work_mgr_sock = false;
        rz.send(()).unwrap();
        loop {
            {
                let timeout = self.poll_timeout();
                let mut items = [self.hb_sock.as_poll_item(1),
                                 self.rq_sock.as_poll_item(1),
                                 self.work_mgr_sock.as_poll_item(1)];
                // Poll until timeout or message is received. Checking for the zmq::POLLIN flag on
                // a poll item's revents will let you know if you have received a message or not
                // on that socket.
                try!(zmq::poll(&mut items, timeout));
                if (items[0].get_revents() & zmq::POLLIN) > 0 {
                    hb_sock = true;
                }
                if (items[1].get_revents() & zmq::POLLIN) > 0 {
                    rq_sock = true;
                }
                if (items[2].get_revents() & zmq::POLLIN) > 0 {
                    work_mgr_sock = true;
                }
            }
            if hb_sock {
                try!(self.process_heartbeat());
                hb_sock = false;
            }
            self.expire_workers();
            if rq_sock {
                try!(self.process_job_status());
                rq_sock = false;
            }
            if work_mgr_sock {
                try!(self.distribute_work());
                work_mgr_sock = false;
            }
        }
    }

    fn poll_timeout(&self) -> i64 {
        if let Some((_, expiry)) = self.workers.front() {
            let timeout = *expiry - Instant::now();
            (timeout.as_secs() as i64 * 1000) + (timeout.subsec_nanos() as i64 / 1000 / 1000)
        } else {
            -1
        }
    }

    fn distribute_work(&mut self) -> Result<()> {
        loop {
            // Take one job from the pending list
            let mut jobs = self.datastore.pending_jobs(1)?;
            // 0 means there are no pending jobs, so we should consume our notice that we have work
            if jobs.len() == 0 {
                try!(self.work_mgr_sock.recv(&mut self.msg, 0));
                break;
            }
            // This unwrap is fine, because we just checked our length
            let mut job = jobs.pop().unwrap();

            match self.workers.pop_front() {
                Some((worker, _)) => {
                    debug!("sending work, worker={:?}, job={:?}", worker, job);
                    if self.rq_sock.send_str(&worker, zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.set_job_state(&job)?;
                        continue;
                    }
                    if self.rq_sock.send(&[], zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.set_job_state(&job)?;
                        continue;
                    }
                    if self.rq_sock.send(&job.write_to_bytes().unwrap(), 0).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.set_job_state(&job)?;
                        continue;
                    }
                }
                None => break,
            }
        }
        Ok(())
    }

    fn expire_workers(&mut self) {
        let now = Instant::now();
        loop {
            if let Some((_, expiry)) = self.workers.front() {
                if expiry >= &now {
                    break;
                }
            } else {
                break;
            }
            let worker = self.workers.pop_front();
            debug!("expiring worker due to inactivity, worker={:?}", worker);
        }
    }

    fn process_heartbeat(&mut self) -> Result<()> {
        try!(self.hb_sock.recv(&mut self.msg, 0));
        let heartbeat: jobsrv::Heartbeat = try!(parse_from_bytes(&self.msg));
        debug!("heartbeat={:?}", heartbeat);
        match heartbeat.get_state() {
            jobsrv::WorkerState::Ready => {
                let now = Instant::now();
                let expiry = now + Duration::from_millis(WORKER_TIMEOUT_MS);
                self.workers.insert(heartbeat.get_endpoint().to_string(), expiry);
            }
            jobsrv::WorkerState::Busy => {
                self.workers.remove(heartbeat.get_endpoint());
            }
        }
        Ok(())
    }

    fn process_job_status(&mut self) -> Result<()> {
        // Pop message delimiter
        try!(self.rq_sock.recv(&mut self.msg, 0));
        // Pop message body
        try!(self.rq_sock.recv(&mut self.msg, 0));
        let job: jobsrv::Job = try!(parse_from_bytes(&self.msg));
        debug!("job_status={:?}", job);
        try!(self.datastore.set_job_state(&job));
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Traits giving structural informations on linear algebra objects or the space they live in.

use std::{f32, f64, i8, i16, i32, i64, u8, u16, u32, u64, isize, usize};
use std::slice::{Iter, IterMut};
use std::ops::{Add, Sub, Mul, Div, Rem, Index, IndexMut, Neg};
use num::{Float, Zero, One};
use traits::operations::{RMul, LMul, Axpy, Transpose, Inv, Absolute};
use traits::geometry::{Dot, Norm, Orig};

/// Basic integral numeric trait.
pub trait BaseNum: Copy + Zero + One +
                   Add<Self, Output = Self> + Sub<Self, Output = Self> +
                   Mul<Self, Output = Self> + Div<Self, Output = Self> +
                   Rem<Self, Output = Self> + PartialEq +
                   Absolute<Self> + Axpy<Self> {
}

/// Basic floating-point number numeric trait.
pub trait BaseFloat: Float + Cast<f64> + BaseNum + Neg {
    /// Archimedes' constant.
    fn pi() -> Self;
    /// 2.0 * pi.
    fn two_pi() -> Self;
    /// pi / 2.0.
    fn frac_pi_2() -> Self;
    /// pi / 3.0.
    fn frac_pi_3() -> Self;
    /// pi / 4.0.
    fn frac_pi_4() -> Self;
    /// pi / 6.0.
    fn frac_pi_6() -> Self;
    /// pi / 8.0.
    fn frac_pi_8() -> Self;
    /// 1.0 / pi.
    fn frac_1_pi() -> Self;
    /// 2.0 / pi.
    fn frac_2_pi() -> Self;
    /// 2.0 / sqrt(pi).
    fn frac_2_sqrt_pi() -> Self;

    /// Euler's number.
    fn e() -> Self;
    /// log2(e).
    fn log2_e() -> Self;
    /// log10(e).
    fn log10_e() -> Self;
    /// ln(2.0).
    fn ln_2() -> Self;
    /// ln(10.0).
    fn ln_10() -> Self;
}

/// Traits of objects which can be created from an object of type `T`.
pub trait Cast<T> {
    /// Converts an element of type `T` to an element of type `Self`.
    fn from(t: T) -> Self;
}

/// Trait of matrices.
///
/// A matrix has rows and columns and are able to multiply them.
pub trait Mat<N, R, C>: Row<R> + Col<C> + RMul<R> + LMul<C> + Index<(usize, usize), Output = N> { }

impl<N, M, R, C> Mat<N, R, C> for M
    where M: Row<R> + Col<C> + RMul<R> + LMul<C> + Index<(usize, usize), Output = N> {
}

/// Trait implemented by square matrices.
pub trait SquareMat<N, V>: Mat<N, V, V> +
                           Mul<Self, Output = Self> + Eye + Transpose + Diag<V> + Inv + Dim + One {
}

impl<N, V, M> SquareMat<N, V> for M
    where M: Mat<N, V, V> + Mul<M, Output = M> + Eye + Transpose + Diag<V> + Inv + Dim + One {
}

/// Trait for constructing the identity matrix
pub trait Eye {
    /// Return the identity matrix of specified dimension
    fn new_identity(dim: usize) -> Self;
}

/// Types that have maximum and minimum value.
pub trait Bounded {
    /// The minimum value.
    #[inline]
    fn min_value() -> Self;
    /// The maximum value.
    #[inline]
    fn max_value() -> Self;
}

// FIXME: return an iterator instead
/// Traits of objects which can form a basis (typically vectors).
pub trait Basis {
    /// Iterates through the canonical basis of the space in which this object lives.
    fn canonical_basis<F: FnMut(Self) -> bool>(F);

    /// Iterates through a basis of the subspace orthogonal to `self`.
    fn orthonormal_subspace_basis<F: FnMut(Self) -> bool>(&Self, F);

    /// Gets the ith element of the canonical basis.
    fn canonical_basis_element(i: usize) -> Option<Self>;
}

/// Trait to access rows of a matrix or a vector.
pub trait Row<R> {
    /// The number of column of `self`.
    fn nrows(&self) -> usize;
    /// Reads the `i`-th row of `self`.
    fn row(&self, i: usize) -> R;
    /// Writes the `i`-th row of `self`.
    fn set_row(&mut self, i: usize, R);

    // FIXME: add iterators on rows: this could be a very good way to generalize _and_ optimize
    // a lot of operations.
}

/// Trait to access columns of a matrix or vector.
pub trait Col<C> {
    /// The number of column of this matrix or vector.
    fn ncols(&self) -> usize;

    /// Reads the `i`-th column of `self`.
    fn col(&self, i: usize) -> C;

    /// Writes the `i`-th column of `self`.
    fn set_col(&mut self, i: usize, C);

    // FIXME: add iterators on columns: this could be a very good way to generalize _and_ optimize
    // a lot of operations.
}

/// Trait to access part of a column of a matrix
pub trait ColSlice<C> {
    /// Returns a view to a slice of a column of a matrix.
    fn col_slice(&self, col_id: usize, row_start: usize, row_end: usize) -> C;
}

/// Trait to access part of a row of a matrix
pub trait RowSlice<R> {
    /// Returns a view to a slice of a row of a matrix.
    fn row_slice(&self, row_id: usize, col_start: usize, col_end: usize) -> R;
}

/// Trait of objects having a spacial dimension known at compile time.
pub trait Dim {
    /// The dimension of the object.
    fn dim(unused_mut: Option<Self>) -> usize;
}

/// Trait to get the diagonal of square matrices.
pub trait Diag<V> {
    /// Creates a new matrix with the given diagonal.
    fn from_diag(diag: &V) -> Self;

    /// Sets the diagonal of this matrix.
    fn set_diag(&mut self, diag: &V);

    /// The diagonal of this matrix.
    fn diag(&self) -> V;
}

/// The shape of an indexable object.
pub trait Shape<I>: Index<I> {
    /// Returns the shape of an indexable object.
    fn shape(&self) -> I;
}

/// This is a workaround of current Rust limitations.
///
/// It exists because the `I` trait cannot be used to express write access.
/// Thus, this is the same as the `I` trait but without the syntactic sugar and with a method
/// to write to a specific index.
pub trait Indexable<I, N>: Shape<I> + IndexMut<I, Output = N> {
    /// Swaps the `i`-th element of `self` with its `j`-th element.
    fn swap(&mut self, i: I, j: I);

    /// Reads the `i`-th element of `self`.
    ///
    /// `i` is not checked.
    unsafe fn unsafe_at(&self, i: I) -> N;
    /// Writes to the `i`-th element of `self`.
    ///
    /// `i` is not checked.
    unsafe fn unsafe_set(&mut self, i: I, N);
}

/// This is a workaround of current Rust limitations.
///
/// Traits of objects which can be iterated through like a vector.
pub trait Iterable<N> {
    /// Gets a vector-like read-only iterator.
    fn iter<'l>(&'l self) -> Iter<'l, N>;
}

/// This is a workaround of current Rust limitations.
///
/// Traits of mutable objects which can be iterated through like a vector.
pub trait IterableMut<N> {
    /// Gets a vector-like read-write iterator.
    fn iter_mut<'l>(&'l mut self) -> IterMut<'l, N>;
}

/*
 * Vec related traits.
 */
/// Trait grouping most common operations on vectors.
pub trait NumVec<N>: Dim +
                     Sub<Self, Output = Self> + Add<Self, Output = Self> +
                     Mul<N, Output = Self> + Div<N, Output = Self> +
                     Index<usize, Output = N> +
                     Zero + PartialEq + Dot<N> + Axpy<N> {
}

/// Trait of vector with components implementing the `BaseFloat` trait.
pub trait FloatVec<N: BaseFloat>: NumVec<N> + Norm<N> + Basis {
}

/*
 * Pnt related traits.
 */
/// Trait that relates a point of an affine space to a vector of the associated vector space.
pub trait PntAsVec<V> {
    /// Converts this point to its associated vector.
    fn to_vec(self) -> V;

    /// Converts a reference to this point to a reference to its associated vector.
    fn as_vec<'a>(&'a self) -> &'a V;

    // NOTE: this is used in some places to overcome some limitations untill the trait reform is
    // done on rustc.
    /// Sets the coordinates of this point to match those of a given vector.
    fn set_coords(&mut self, coords: V);
}

/// Trait grouping most common operations on points.
// XXX: the vector space element `V` should be an associated type. Though this would prevent V from
// having bounds (they are not supported yet). So, for now, we will just use a type parameter.
pub trait NumPnt<N, V>:
          Copy +
          PntAsVec<V> +
          Dim +
          Orig +
          PartialEq +
          Axpy<N> +
          Sub<Self, Output = V> +
          Mul<N, Output = Self> +
          Div<N, Output = Self> +
          Add<V, Output = Self> +
          Index<usize, Output = N> { // FIXME: + Sub<V, Self>
}

/// Trait of points with components implementing the `BaseFloat` trait.
pub trait FloatPnt<N: BaseFloat, V: Norm<N>>: NumPnt<N, V> + Sized {
    /// Computes the square distance between two points.
    #[inline]
    fn sqdist(&self, other: &Self) -> N {
        (*self - *other).sqnorm()
    }

    /// Computes the distance between two points.
    #[inline]
    fn dist(&self, other: &Self) -> N {
        (*self - *other).norm()
    }
}

/*
 *
 *
 * Some implementations for builtin types.
 *
 *
 */
// Bounded
macro_rules! impl_bounded(
    ($n: ty, $min: expr, $max: expr) => {
        impl Bounded for $n {
            #[inline]
            fn min_value() -> $n {
                $min
            }

            #[inline]
            fn max_value() -> $n {
                $max
            }
        }
    }
);

impl_bounded!(f32, f32::MIN, f32::MAX);
impl_bounded!(f64, f64::MIN, f64::MAX);
impl_bounded!(i8, i8::MIN, i8::MAX);
impl_bounded!(i16, i16::MIN, i16::MAX);
impl_bounded!(i32, i32::MIN, i32::MAX);
impl_bounded!(i64, i64::MIN, i64::MAX);
impl_bounded!(isize, isize::MIN, isize::MAX);
impl_bounded!(u8, u8::MIN, u8::MAX);
impl_bounded!(u16, u16::MIN, u16::MAX);
impl_bounded!(u32, u32::MIN, u32::MAX);
impl_bounded!(u64, u64::MIN, u64::MAX);
impl_bounded!(usize, usize::MIN, usize::MAX);

// BaseFloat
macro_rules! impl_base_float(
    ($n: ident) => {
        impl BaseFloat for $n {
            /// Archimedes' constant.
            fn pi() -> $n {
                $n::consts::PI
            }

            /// 2.0 * pi.
            fn two_pi() -> $n {
                2.0 * $n::consts::PI
            }

            /// pi / 2.0.
            fn frac_pi_2() -> $n {
                $n::consts::FRAC_PI_2
            }

            /// pi / 3.0.
            fn frac_pi_3() -> $n {
                $n::consts::FRAC_PI_3
            }

            /// pi / 4.0.
            fn frac_pi_4() -> $n {
                $n::consts::FRAC_PI_4
            }

            /// pi / 6.0.
            fn frac_pi_6() -> $n {
                $n::consts::FRAC_PI_6
            }

            /// pi / 8.0.
            fn frac_pi_8() -> $n {
                $n::consts::FRAC_PI_8
            }

            /// 1.0 / pi.
            fn frac_1_pi() -> $n {
                $n::consts::FRAC_1_PI
            }

            /// 2.0 / pi.
            fn frac_2_pi() -> $n {
                $n::consts::FRAC_2_PI
            }

            /// 2.0 / sqrt(pi).
            fn frac_2_sqrt_pi() -> $n {
                $n::consts::FRAC_2_SQRT_PI
            }

            /// Euler's number.
            fn e() -> $n {
                $n::consts::E
            }

            /// log2(e).
            fn log2_e() -> $n {
                $n::consts::LOG2_E
            }

            /// log10(e).
            fn log10_e() -> $n {
                $n::consts::LOG10_E
            }

            /// ln(2.0).
            fn ln_2() -> $n {
                $n::consts::LN_2
            }

            /// ln(10.0).
            fn ln_10() -> $n {
                $n::consts::LN_10
            }
        }
    }
);

impl BaseNum for i8 { }
impl BaseNum for i16 { }
impl BaseNum for i32 { }
impl BaseNum for i64 { }
impl BaseNum for isize { }
impl BaseNum for u8 { }
impl BaseNum for u16 { }
impl BaseNum for u32 { }
impl BaseNum for u64 { }
impl BaseNum for usize { }
impl BaseNum for f32 { }
impl BaseNum for f64 { }

impl_base_float!(f32);
impl_base_float!(f64);
<commit_msg>Qualify the output of FloatVec Neg<commit_after>//! Traits giving structural informations on linear algebra objects or the space they live in.

use std::{f32, f64, i8, i16, i32, i64, u8, u16, u32, u64, isize, usize};
use std::slice::{Iter, IterMut};
use std::ops::{Add, Sub, Mul, Div, Rem, Index, IndexMut, Neg};
use num::{Float, Zero, One};
use traits::operations::{RMul, LMul, Axpy, Transpose, Inv, Absolute};
use traits::geometry::{Dot, Norm, Orig};

/// Basic integral numeric trait.
pub trait BaseNum: Copy + Zero + One +
                   Add<Self, Output = Self> + Sub<Self, Output = Self> +
                   Mul<Self, Output = Self> + Div<Self, Output = Self> +
                   Rem<Self, Output = Self> + PartialEq +
                   Absolute<Self> + Axpy<Self> {
}

/// Basic floating-point number numeric trait.
pub trait BaseFloat: Float + Cast<f64> + BaseNum + Neg {
    /// Archimedes' constant.
    fn pi() -> Self;
    /// 2.0 * pi.
    fn two_pi() -> Self;
    /// pi / 2.0.
    fn frac_pi_2() -> Self;
    /// pi / 3.0.
    fn frac_pi_3() -> Self;
    /// pi / 4.0.
    fn frac_pi_4() -> Self;
    /// pi / 6.0.
    fn frac_pi_6() -> Self;
    /// pi / 8.0.
    fn frac_pi_8() -> Self;
    /// 1.0 / pi.
    fn frac_1_pi() -> Self;
    /// 2.0 / pi.
    fn frac_2_pi() -> Self;
    /// 2.0 / sqrt(pi).
    fn frac_2_sqrt_pi() -> Self;

    /// Euler's number.
    fn e() -> Self;
    /// log2(e).
    fn log2_e() -> Self;
    /// log10(e).
    fn log10_e() -> Self;
    /// ln(2.0).
    fn ln_2() -> Self;
    /// ln(10.0).
    fn ln_10() -> Self;
}

/// Traits of objects which can be created from an object of type `T`.
pub trait Cast<T> {
    /// Converts an element of type `T` to an element of type `Self`.
    fn from(t: T) -> Self;
}

/// Trait of matrices.
///
/// A matrix has rows and columns and are able to multiply them.
pub trait Mat<N, R, C>: Row<R> + Col<C> + RMul<R> + LMul<C> + Index<(usize, usize), Output = N> { }

impl<N, M, R, C> Mat<N, R, C> for M
    where M: Row<R> + Col<C> + RMul<R> + LMul<C> + Index<(usize, usize), Output = N> {
}

/// Trait implemented by square matrices.
pub trait SquareMat<N, V>: Mat<N, V, V> +
                           Mul<Self, Output = Self> + Eye + Transpose + Diag<V> + Inv + Dim + One {
}

impl<N, V, M> SquareMat<N, V> for M
    where M: Mat<N, V, V> + Mul<M, Output = M> + Eye + Transpose + Diag<V> + Inv + Dim + One {
}

/// Trait for constructing the identity matrix
pub trait Eye {
    /// Return the identity matrix of specified dimension
    fn new_identity(dim: usize) -> Self;
}

/// Types that have maximum and minimum value.
pub trait Bounded {
    /// The minimum value.
    #[inline]
    fn min_value() -> Self;
    /// The maximum value.
    #[inline]
    fn max_value() -> Self;
}

// FIXME: return an iterator instead
/// Traits of objects which can form a basis (typically vectors).
pub trait Basis {
    /// Iterates through the canonical basis of the space in which this object lives.
    fn canonical_basis<F: FnMut(Self) -> bool>(F);

    /// Iterates through a basis of the subspace orthogonal to `self`.
    fn orthonormal_subspace_basis<F: FnMut(Self) -> bool>(&Self, F);

    /// Gets the ith element of the canonical basis.
    fn canonical_basis_element(i: usize) -> Option<Self>;
}

/// Trait to access rows of a matrix or a vector.
pub trait Row<R> {
    /// The number of column of `self`.
    fn nrows(&self) -> usize;
    /// Reads the `i`-th row of `self`.
    fn row(&self, i: usize) -> R;
    /// Writes the `i`-th row of `self`.
    fn set_row(&mut self, i: usize, R);

    // FIXME: add iterators on rows: this could be a very good way to generalize _and_ optimize
    // a lot of operations.
}

/// Trait to access columns of a matrix or vector.
pub trait Col<C> {
    /// The number of column of this matrix or vector.
    fn ncols(&self) -> usize;

    /// Reads the `i`-th column of `self`.
    fn col(&self, i: usize) -> C;

    /// Writes the `i`-th column of `self`.
    fn set_col(&mut self, i: usize, C);

    // FIXME: add iterators on columns: this could be a very good way to generalize _and_ optimize
    // a lot of operations.
}

/// Trait to access part of a column of a matrix
pub trait ColSlice<C> {
    /// Returns a view to a slice of a column of a matrix.
    fn col_slice(&self, col_id: usize, row_start: usize, row_end: usize) -> C;
}

/// Trait to access part of a row of a matrix
pub trait RowSlice<R> {
    /// Returns a view to a slice of a row of a matrix.
    fn row_slice(&self, row_id: usize, col_start: usize, col_end: usize) -> R;
}

/// Trait of objects having a spacial dimension known at compile time.
pub trait Dim {
    /// The dimension of the object.
    fn dim(unused_mut: Option<Self>) -> usize;
}

/// Trait to get the diagonal of square matrices.
pub trait Diag<V> {
    /// Creates a new matrix with the given diagonal.
    fn from_diag(diag: &V) -> Self;

    /// Sets the diagonal of this matrix.
    fn set_diag(&mut self, diag: &V);

    /// The diagonal of this matrix.
    fn diag(&self) -> V;
}

/// The shape of an indexable object.
pub trait Shape<I>: Index<I> {
    /// Returns the shape of an indexable object.
    fn shape(&self) -> I;
}

/// This is a workaround of current Rust limitations.
///
/// It exists because the `I` trait cannot be used to express write access.
/// Thus, this is the same as the `I` trait but without the syntactic sugar and with a method
/// to write to a specific index.
pub trait Indexable<I, N>: Shape<I> + IndexMut<I, Output = N> {
    /// Swaps the `i`-th element of `self` with its `j`-th element.
    fn swap(&mut self, i: I, j: I);

    /// Reads the `i`-th element of `self`.
    ///
    /// `i` is not checked.
    unsafe fn unsafe_at(&self, i: I) -> N;
    /// Writes to the `i`-th element of `self`.
    ///
    /// `i` is not checked.
    unsafe fn unsafe_set(&mut self, i: I, N);
}

/// This is a workaround of current Rust limitations.
///
/// Traits of objects which can be iterated through like a vector.
pub trait Iterable<N> {
    /// Gets a vector-like read-only iterator.
    fn iter<'l>(&'l self) -> Iter<'l, N>;
}

/// This is a workaround of current Rust limitations.
///
/// Traits of mutable objects which can be iterated through like a vector.
pub trait IterableMut<N> {
    /// Gets a vector-like read-write iterator.
    fn iter_mut<'l>(&'l mut self) -> IterMut<'l, N>;
}

/*
 * Vec related traits.
 */
/// Trait grouping most common operations on vectors.
pub trait NumVec<N>: Dim +
                     Sub<Self, Output = Self> + Add<Self, Output = Self> +
                     Mul<N, Output = Self> + Div<N, Output = Self> +
                     Index<usize, Output = N> +
                     Zero + PartialEq + Dot<N> + Axpy<N> {
}

/// Trait of vector with components implementing the `BaseFloat` trait.
pub trait FloatVec<N: BaseFloat>: NumVec<N> + Norm<N> + Neg<Output = Self> + Basis {
}

/*
 * Pnt related traits.
 */
/// Trait that relates a point of an affine space to a vector of the associated vector space.
pub trait PntAsVec<V> {
    /// Converts this point to its associated vector.
    fn to_vec(self) -> V;

    /// Converts a reference to this point to a reference to its associated vector.
    fn as_vec<'a>(&'a self) -> &'a V;

    // NOTE: this is used in some places to overcome some limitations untill the trait reform is
    // done on rustc.
    /// Sets the coordinates of this point to match those of a given vector.
    fn set_coords(&mut self, coords: V);
}

/// Trait grouping most common operations on points.
// XXX: the vector space element `V` should be an associated type. Though this would prevent V from
// having bounds (they are not supported yet). So, for now, we will just use a type parameter.
pub trait NumPnt<N, V>:
          Copy +
          PntAsVec<V> +
          Dim +
          Orig +
          PartialEq +
          Axpy<N> +
          Sub<Self, Output = V> +
          Mul<N, Output = Self> +
          Div<N, Output = Self> +
          Add<V, Output = Self> +
          Index<usize, Output = N> { // FIXME: + Sub<V, Self>
}

/// Trait of points with components implementing the `BaseFloat` trait.
pub trait FloatPnt<N: BaseFloat, V: Norm<N>>: NumPnt<N, V> + Sized {
    /// Computes the square distance between two points.
    #[inline]
    fn sqdist(&self, other: &Self) -> N {
        (*self - *other).sqnorm()
    }

    /// Computes the distance between two points.
    #[inline]
    fn dist(&self, other: &Self) -> N {
        (*self - *other).norm()
    }
}

/*
 *
 *
 * Some implementations for builtin types.
 *
 *
 */
// Bounded
macro_rules! impl_bounded(
    ($n: ty, $min: expr, $max: expr) => {
        impl Bounded for $n {
            #[inline]
            fn min_value() -> $n {
                $min
            }

            #[inline]
            fn max_value() -> $n {
                $max
            }
        }
    }
);

impl_bounded!(f32, f32::MIN, f32::MAX);
impl_bounded!(f64, f64::MIN, f64::MAX);
impl_bounded!(i8, i8::MIN, i8::MAX);
impl_bounded!(i16, i16::MIN, i16::MAX);
impl_bounded!(i32, i32::MIN, i32::MAX);
impl_bounded!(i64, i64::MIN, i64::MAX);
impl_bounded!(isize, isize::MIN, isize::MAX);
impl_bounded!(u8, u8::MIN, u8::MAX);
impl_bounded!(u16, u16::MIN, u16::MAX);
impl_bounded!(u32, u32::MIN, u32::MAX);
impl_bounded!(u64, u64::MIN, u64::MAX);
impl_bounded!(usize, usize::MIN, usize::MAX);

// BaseFloat
macro_rules! impl_base_float(
    ($n: ident) => {
        impl BaseFloat for $n {
            /// Archimedes' constant.
            fn pi() -> $n {
                $n::consts::PI
            }

            /// 2.0 * pi.
            fn two_pi() -> $n {
                2.0 * $n::consts::PI
            }

            /// pi / 2.0.
            fn frac_pi_2() -> $n {
                $n::consts::FRAC_PI_2
            }

            /// pi / 3.0.
            fn frac_pi_3() -> $n {
                $n::consts::FRAC_PI_3
            }

            /// pi / 4.0.
            fn frac_pi_4() -> $n {
                $n::consts::FRAC_PI_4
            }

            /// pi / 6.0.
            fn frac_pi_6() -> $n {
                $n::consts::FRAC_PI_6
            }

            /// pi / 8.0.
            fn frac_pi_8() -> $n {
                $n::consts::FRAC_PI_8
            }

            /// 1.0 / pi.
            fn frac_1_pi() -> $n {
                $n::consts::FRAC_1_PI
            }

            /// 2.0 / pi.
            fn frac_2_pi() -> $n {
                $n::consts::FRAC_2_PI
            }

            /// 2.0 / sqrt(pi).
            fn frac_2_sqrt_pi() -> $n {
                $n::consts::FRAC_2_SQRT_PI
            }

            /// Euler's number.
            fn e() -> $n {
                $n::consts::E
            }

            /// log2(e).
            fn log2_e() -> $n {
                $n::consts::LOG2_E
            }

            /// log10(e).
            fn log10_e() -> $n {
                $n::consts::LOG10_E
            }

            /// ln(2.0).
            fn ln_2() -> $n {
                $n::consts::LN_2
            }

            /// ln(10.0).
            fn ln_10() -> $n {
                $n::consts::LN_10
            }
        }
    }
);

impl BaseNum for i8 { }
impl BaseNum for i16 { }
impl BaseNum for i32 { }
impl BaseNum for i64 { }
impl BaseNum for isize { }
impl BaseNum for u8 { }
impl BaseNum for u16 { }
impl BaseNum for u32 { }
impl BaseNum for u64 { }
impl BaseNum for usize { }
impl BaseNum for f32 { }
impl BaseNum for f64 { }

impl_base_float!(f32);
impl_base_float!(f64);
<|endoftext|>


<commit_before>use std::cell::{Cell, RefCell};
use std::collections::HashSet;
use std::fmt;
use std::io::{BufWriter, Bytes, Result, Write};

use super::nodes::{self, *};
use super::typ::*;

macro_rules! iwrite {
    ($self: ident, $($arg:tt)*) => (
        $self.writeln(format_args!($($arg)*))
    );
}

macro_rules! lnwrite {
    ($self: ident, $($arg:tt)*) => (
        $self.writeln(format_args!($($arg)*))
    );
}

pub struct PrinterOptions {
    pub print_pointers: bool,
}

impl Default for PrinterOptions {
    fn default() -> Self {
        Self {
            print_pointers: false,
        }
    }
}

pub struct Printer<O: Write> {
    output: RefCell<O>,
    indent: Cell<u8>,
    print_pointers: bool,
}

impl<O: Write> Printer<O> {
    pub fn new(output: O) -> Self {
        Self::new_with_options(output, PrinterOptions::default())
    }

    pub fn new_with_options(output: O, options: PrinterOptions) -> Self {
        Self {
            output: RefCell::new(output),
            indent: Cell::new(0),
            print_pointers: options.print_pointers,
        }
    }

    fn indented<F: FnOnce(&Self) -> Result<()>>(&self, inner: F) -> Result<()> {
        self.indented_steps(1, inner)
    }

    fn indented_steps<F: FnOnce(&Self) -> Result<()>>(&self, steps: u8, inner: F) -> Result<()> {
        let previous = self.indent.take();
        self.indent.set(previous + (2 * steps));
        let result = inner(self);
        self.indent.set(previous);
        result
    }

    pub fn print_module(&self, module: Module) -> Result<()> {
        for (index, statement) in module.statements.iter().enumerate() {
            use ModuleStatement::*;
            match statement {
                Func(func) => self.print_func(func, index == 0)?,
            }
        }
        self.lnwrite("")?;
        Ok(())
    }

    fn print_func(&self, func: &nodes::Func, first: bool) -> Result<()> {
        let opener = format!("func {}(", func.name);
        if first {
            self.iwrite(opener)?;
        } else {
            self.lnwrite(opener)?;
        }
        if !func.arguments.is_empty() {
            self.write("\n")?;
            self.indented(|this| {
                for argument in func.arguments.iter() {
                    this.iwrite(format!("{}: ", argument.name))?;
                    this.write_type(&argument.typ, true)?;
                    this.write(",\n")?;
                }
                Ok(())
            })?;
        }
        self.iwrite("): ")?;
        let retrn = match &func.typ {
            Type::Func(func) => func.retrn.borrow(),
            other @ _ => unreachable!("Func node has non-Func type: {:?}", other),
        };
        self.write_type(&retrn, false)?;
        self.write(" ")?;
        match &func.body {
            FuncBody::Block(block) => self.print_block(block, false),
        }
    }

    fn write_type(&self, typ: &Type, with_children: bool) -> Result<()> {
        self.write_recursive_type(typ, with_children, &mut HashSet::new())
    }

    // Write type recursively with guards against infinite recursion.
    fn write_recursive_type(
        &self,
        typ: &Type,
        with_children: bool,
        tracker: &mut HashSet<usize>,
    ) -> Result<()> {
        // True if we're writing a type we've written before.
        let recursive = !tracker.insert(typ.id());

        match typ {
            Type::Func(func) => {
                self.write(format!(
                    "{}#{}(",
                    func.name.clone().unwrap_or("".to_string()),
                    func.id
                ))?;
                let arguments = func.arguments.borrow();
                if !arguments.is_empty() {
                    if with_children {
                        for argument in arguments.iter() {
                            self.indented(|this1| {
                                this1.lnwrite("")?;
                                this1.write_recursive_type(argument, true, tracker)?;
                                this1.write(",")
                            })?;
                        }
                        self.lnwrite(")")?;
                    } else {
                        self.write(format!("{})", arguments.len()))?;
                    }
                } else {
                    self.write(")")?;
                }
                self.write_pointer(&**func)
            }
            Type::Object(object) => self.write(format!("{}", object.class.name())),
            Type::Generic(outer) => {
                let generic = outer.borrow();
                self.write(format!("${}", generic.id))?;
                if with_children && !recursive {
                    self.write("(")?;
                    self.indented(|this| this.write_constraints(&generic.constraints, tracker))?;
                    self.write(")")?;
                }
                self.write_pointer(&**outer)?;
                Ok(())
            }
            Type::Tuple(tuple) => self.write(format!("({})", tuple.members.len())),
            Type::Variable(variable) => {
                let variable = &*variable.borrow();
                match variable {
                    Variable::Substitute { substitute, .. } => {
                        self.write("S(")?;
                        if recursive {
                            self.write("...")?;
                        } else {
                            self.write_type(&*substitute, with_children)?;
                        }
                        self.write(")")?;
                    }
                    Variable::Unbound { id, .. } => {
                        self.write(format!("U({})", id))?;
                    }
                    Variable::Generic { generic, .. } => {
                        self.write("G(")?;
                        self.write(format!("{}", generic.id))?;
                        if with_children && !recursive {
                            self.indented(|this| {
                                this.write_constraints(&generic.constraints, tracker)
                            })?;
                        }
                        self.write(")")?;
                    }
                }
                self.write_pointer(variable)
            }
            _ => unreachable!("Cannot print type: {:?}", typ),
        }
    }

    fn write_constraints(
        &self,
        constraints: &Vec<GenericConstraint>,
        tracker: &mut HashSet<usize>,
    ) -> Result<()> {
        if constraints.is_empty() {
            return Ok(());
        }
        self.lnwrite("where")?;
        self.indented(|this1| {
            for constraint in constraints {
                use GenericConstraint::*;
                match constraint {
                    Property(property) => {
                        this1.lnwrite(format!("{}: ", property.name))?;
                        this1.write_recursive_type(&property.typ, true, tracker)?;
                    }
                    Callable(callable) => {
                        if callable.arguments.is_empty() {
                            this1.lnwrite("(): ")?;
                        } else {
                            this1.lnwrite("(\n")?;
                            for argument in callable.arguments.iter() {
                                this1.indented(|this2| {
                                    this2.iwrite("")?;
                                    this2.write_recursive_type(argument, true, tracker)?;
                                    this2.write(",")
                                })?;
                            }
                            this1.lnwrite("): ")?;
                        }
                        this1.write_recursive_type(&callable.retrn, true, tracker)?;
                    }
                }
            }
            Ok(())
        })
    }

    fn print_var(&self, var: &Var, first: bool) -> Result<()> {
        let opener = format!("var {}: ", &var.name.name);
        if first {
            self.iwrite(opener)?;
        } else {
            self.lnwrite(opener)?;
        }
        self.write_type(&var.typ, true)?;
        if let Some(initializer) = &var.initializer {
            self.indented(|this| {
                this.write(" =")?;
                this.print_expression(initializer)
            })?;
        }
        Ok(())
    }

    fn print_block(&self, block: &Block, initial_indent: bool) -> Result<()> {
        let opener = "{";
        if initial_indent {
            self.iwrite(opener)?;
        } else {
            self.write(opener)?;
        }

        self.indented(|this| {
            for (_, statement) in block.statements.iter().enumerate() {
                use BlockStatement::*;
                match statement {
                    Expression(expression) => {
                        self.print_expression(expression)?;
                    }
                    Func(func) => self.print_func(func, false)?,
                    Var(var) => self.print_var(var, false)?,
                }
            }
            Ok(())
        })?;
        self.write("\n")?;
        self.iwrite("}")
    }

    fn print_expression(&self, expression: &Expression) -> Result<()> {
        use Expression::*;
        match expression {
            Closure(closure) => self.print_closure(closure),
            Identifier(identifier) => self.print_identifier(identifier),
            Infix(infix) => self.print_infix(infix),
            LiteralInt(literal) => self.lnwrite(format!("{}", literal.value)),
            PostfixCall(call) => self.print_postfix_call(call, 0).map(|_| ()),
            PostfixProperty(property) => self.print_postfix_property(property, 0).map(|_| ()),
        }
    }

    fn print_closure(&self, closure: &Closure) -> Result<()> {
        self.lnwrite("Closure(")?;
        if !closure.arguments.is_empty() {
            self.write("\n")?;
            self.indented(|this| {
                for argument in closure.arguments.iter() {
                    this.iwrite(format!("{}: ", argument.name))?;
                    this.write_type(&argument.typ, true)?;
                    this.write(",\n")?;
                }
                Ok(())
            })?;
        }
        self.iwrite("): ")?;
        let retrn = match &closure.typ {
            Type::Func(func) => func.retrn.borrow(),
            other @ _ => unreachable!("Closure node has non-Func type: {:?}", other),
        };
        self.write_type(&retrn, false)?;
        self.write(" ")?;
        match &*closure.body {
            ClosureBody::Block(block) => self.print_block(block, false),
            ClosureBody::Expression(expression) => self.indented(|this| {
                this.write("(")?;
                this.print_expression(expression)?;
                this.write(")")
            }),
        }
    }

    fn print_identifier(&self, identifier: &Identifier) -> Result<()> {
        self.lnwrite(format!("Identifier({}): ", identifier.name.name))?;
        self.write_type(&identifier.typ, false)
    }

    fn print_infix(&self, infix: &Infix) -> Result<()> {
        self.lnwrite("Infix(")?;
        self.indented(|this1| {
            this1.lnwrite("lhs:")?;
            this1.indented(|this2| this2.print_expression(&infix.lhs))?;
            this1.lnwrite(format!("op: {}", infix.op.to_string()))?;
            this1.lnwrite(format!("rhs:"))?;
            this1.indented(|this2| this2.print_expression(&infix.rhs))
        })?;
        self.write(")")
    }

    fn is_postfix(&self, expression: &Expression) -> bool {
        use Expression::*;
        match expression {
            PostfixProperty(_) => true,
            _ => false,
        }
    }

    fn print_postfix_call(&self, call: &PostfixCall, current: u8) -> Result<u8> {
        let max = self.print_postfix_target(&*call.target, current)?;
        // Have a +1 so that we always indent at least one step.
        self.indented_steps(max - current + 1, |this1| {
            this1.lnwrite(format!("Call("))?;
            if !call.arguments.is_empty() {
                for argument in call.arguments.iter() {
                    this1.indented(|this2| {
                        this2.print_expression(argument)?;
                        this2.write(",")
                    })?;
                }
            }
            this1.write("): ")?;
            this1.write_type(&call.typ, false)
        })?;
        Ok(max)
    }

    fn print_postfix_property(&self, property: &PostfixProperty, current: u8) -> Result<u8> {
        let max = self.print_postfix_target(&*property.target, current)?;
        self.indented_steps(max - current + 1, |this| {
            this.lnwrite(format!("Property({}): ", property.property.name))?;
            this.write_type(&property.typ, false)
        })?;
        Ok(max)
    }

    fn print_postfix_target(&self, target: &Expression, current: u8) -> Result<u8> {
        match target {
            // Links in the chain
            Expression::PostfixCall(target) => self.print_postfix_call(target, current + 1),
            Expression::PostfixProperty(target) => self.print_postfix_property(target, current + 1),
            // Tail of the chain
            other @ _ => {
                self.print_expression(other)?;
                Ok(current)
            }
        }
    }

    /// Write a string.
    fn write<S: AsRef<str>>(&self, string: S) -> Result<()> {
        self.write_output(string.as_ref().as_bytes())
    }

    /// Write indentation and then a string.
    fn iwrite<S: AsRef<str>>(&self, string: S) -> Result<()> {
        let indented = " ".repeat(self.indent.get() as usize);
        self.write_output(indented.as_bytes())
            .and_then(|_| self.write_output(string.as_ref().as_bytes()))
            .map(|_| ())
    }

    /// Write a newline, indentation, and then a string.
    fn lnwrite<S: AsRef<str>>(&self, string: S) -> Result<()> {
        let indented = " ".repeat(self.indent.get() as usize);
        self.write_output("\n".as_bytes())
            .and_then(|_| self.write_output(indented.as_bytes()))
            .and_then(|_| self.write_output(string.as_ref().as_bytes()))
    }

    fn write_output(&self, bytes: &[u8]) -> Result<()> {
        let mut output = self.output.borrow_mut();
        output.write(bytes).map(|_| ())
    }

    fn write_pointer<P>(&self, ptr: *const P) -> Result<()> {
        if !self.print_pointers {
            return Ok(());
        }
        self.write(format!(" @ {:p}", ptr))
    }

    // fn write_output(&self, bytes: &[u8]) -> Result<()> {
    //     let result = if let Some(buffer) = &mut *self.buffer.borrow_mut() {
    //         buffer.write(bytes)
    //     } else {
    //         let mut output = self.output.borrow_mut();
    //         output.write(bytes)
    //     };
    //     result.map(|_| ())
    // }
    //
    // fn buffer<F: FnOnce(&Self) -> Result<()>>(&self, inner: F) -> Result<Vec<u8>> {
    //     {
    //         let mut buffer = self.buffer.borrow_mut();
    //         *buffer = Some(vec![]);
    //     }
    //     inner(self)?;
    //     let buffer = self.buffer.replace(None);
    //     Ok(buffer.unwrap())
    // }
}
<commit_msg>Improve printing of function and object types<commit_after>use std::cell::{Cell, RefCell};
use std::collections::HashSet;
use std::fmt;
use std::io::{BufWriter, Bytes, Result, Write};

use super::nodes::{self, *};
use super::typ::*;

macro_rules! iwrite {
    ($self: ident, $($arg:tt)*) => (
        $self.writeln(format_args!($($arg)*))
    );
}

macro_rules! lnwrite {
    ($self: ident, $($arg:tt)*) => (
        $self.writeln(format_args!($($arg)*))
    );
}

pub struct PrinterOptions {
    pub print_pointers: bool,
}

impl Default for PrinterOptions {
    fn default() -> Self {
        Self {
            print_pointers: false,
        }
    }
}

pub struct Printer<O: Write> {
    output: RefCell<O>,
    indent: Cell<u8>,
    print_pointers: bool,
}

impl<O: Write> Printer<O> {
    pub fn new(output: O) -> Self {
        Self::new_with_options(output, PrinterOptions::default())
    }

    pub fn new_with_options(output: O, options: PrinterOptions) -> Self {
        Self {
            output: RefCell::new(output),
            indent: Cell::new(0),
            print_pointers: options.print_pointers,
        }
    }

    fn indented<F: FnOnce(&Self) -> Result<()>>(&self, inner: F) -> Result<()> {
        self.indented_steps(1, inner)
    }

    fn indented_steps<F: FnOnce(&Self) -> Result<()>>(&self, steps: u8, inner: F) -> Result<()> {
        let previous = self.indent.take();
        self.indent.set(previous + (2 * steps));
        let result = inner(self);
        self.indent.set(previous);
        result
    }

    pub fn print_module(&self, module: Module) -> Result<()> {
        for (index, statement) in module.statements.iter().enumerate() {
            use ModuleStatement::*;
            match statement {
                Func(func) => self.print_func(func, index == 0)?,
            }
        }
        self.lnwrite("")?;
        Ok(())
    }

    fn print_func(&self, func: &nodes::Func, first: bool) -> Result<()> {
        let opener = format!("func {}(", func.name);
        if first {
            self.iwrite(opener)?;
        } else {
            self.lnwrite(opener)?;
        }
        if !func.arguments.is_empty() {
            self.write("\n")?;
            self.indented(|this| {
                for argument in func.arguments.iter() {
                    this.iwrite(format!("{}: ", argument.name))?;
                    this.write_type(&argument.typ, true)?;
                    this.write(",\n")?;
                }
                Ok(())
            })?;
        }
        self.iwrite("): ")?;
        let retrn = match &func.typ {
            Type::Func(func) => func.retrn.borrow(),
            other @ _ => unreachable!("Func node has non-Func type: {:?}", other),
        };
        self.write_type(&retrn, false)?;
        self.write(" ")?;
        match &func.body {
            FuncBody::Block(block) => self.print_block(block, false),
        }
    }

    fn write_type(&self, typ: &Type, with_children: bool) -> Result<()> {
        self.write_recursive_type(typ, with_children, &mut HashSet::new())
    }

    // Write type recursively with guards against infinite recursion.
    fn write_recursive_type(
        &self,
        typ: &Type,
        with_children: bool,
        tracker: &mut HashSet<usize>,
    ) -> Result<()> {
        // True if we're writing a type we've written before.
        let recursive = !tracker.insert(typ.id());

        match typ {
            Type::Func(func) => {
                self.write(format!(
                    "{}#{}(",
                    func.name.clone().unwrap_or("".to_string()),
                    func.id
                ))?;
                let arguments = func.arguments.borrow();
                if !arguments.is_empty() {
                    if with_children {
                        for argument in arguments.iter() {
                            self.indented(|this1| {
                                this1.lnwrite("")?;
                                this1.write_recursive_type(argument, true, tracker)?;
                                this1.write(",")
                            })?;
                        }
                        self.lnwrite("): ")?;
                    } else {
                        self.write(format!("{}): ", arguments.len()))?;
                    }
                } else {
                    self.write("): ")?;
                }
                self.write_recursive_type(&func.retrn.borrow(), true, tracker)?;
                self.write_pointer(&**func)
            }
            Type::Object(object) => {
                self.write(format!("{}", object.class.name()))?;
                self.write_pointer(&**object)
            },
            Type::Generic(outer) => {
                let generic = outer.borrow();
                self.write(format!("${}", generic.id))?;
                if with_children && !recursive {
                    self.write("(")?;
                    self.indented(|this| this.write_constraints(&generic.constraints, tracker))?;
                    self.write(")")?;
                }
                self.write_pointer(&**outer)
            }
            Type::Tuple(tuple) => self.write(format!("({})", tuple.members.len())),
            Type::Variable(variable) => {
                let variable = &*variable.borrow();
                match variable {
                    Variable::Substitute { substitute, .. } => {
                        self.write("S(")?;
                        if recursive {
                            self.write("...")?;
                        } else {
                            self.write_type(&*substitute, with_children)?;
                        }
                        self.write(")")?;
                    }
                    Variable::Unbound { id, .. } => {
                        self.write(format!("U({})", id))?;
                    }
                    Variable::Generic { generic, .. } => {
                        self.write("G(")?;
                        self.write(format!("{}", generic.id))?;
                        if with_children && !recursive {
                            self.indented(|this| {
                                this.write_constraints(&generic.constraints, tracker)
                            })?;
                        }
                        self.write(")")?;
                    }
                }
                self.write_pointer(variable)
            }
            _ => unreachable!("Cannot print type: {:?}", typ),
        }
    }

    fn write_constraints(
        &self,
        constraints: &Vec<GenericConstraint>,
        tracker: &mut HashSet<usize>,
    ) -> Result<()> {
        if constraints.is_empty() {
            return Ok(());
        }
        self.lnwrite("where")?;
        self.indented(|this1| {
            for constraint in constraints {
                use GenericConstraint::*;
                match constraint {
                    Property(property) => {
                        this1.lnwrite(format!("{}: ", property.name))?;
                        this1.write_recursive_type(&property.typ, true, tracker)?;
                    }
                    Callable(callable) => {
                        if callable.arguments.is_empty() {
                            this1.lnwrite("(): ")?;
                        } else {
                            this1.lnwrite("(\n")?;
                            for argument in callable.arguments.iter() {
                                this1.indented(|this2| {
                                    this2.iwrite("")?;
                                    this2.write_recursive_type(argument, true, tracker)?;
                                    this2.write(",")
                                })?;
                            }
                            this1.lnwrite("): ")?;
                        }
                        this1.write_recursive_type(&callable.retrn, true, tracker)?;
                    }
                }
            }
            Ok(())
        })
    }

    fn print_var(&self, var: &Var, first: bool) -> Result<()> {
        let opener = format!("var {}: ", &var.name.name);
        if first {
            self.iwrite(opener)?;
        } else {
            self.lnwrite(opener)?;
        }
        self.write_type(&var.typ, true)?;
        if let Some(initializer) = &var.initializer {
            self.indented(|this| {
                this.write(" =")?;
                this.print_expression(initializer)
            })?;
        }
        Ok(())
    }

    fn print_block(&self, block: &Block, initial_indent: bool) -> Result<()> {
        let opener = "{";
        if initial_indent {
            self.iwrite(opener)?;
        } else {
            self.write(opener)?;
        }

        self.indented(|this| {
            for (_, statement) in block.statements.iter().enumerate() {
                use BlockStatement::*;
                match statement {
                    Expression(expression) => {
                        self.print_expression(expression)?;
                    }
                    Func(func) => self.print_func(func, false)?,
                    Var(var) => self.print_var(var, false)?,
                }
            }
            Ok(())
        })?;
        self.write("\n")?;
        self.iwrite("}")
    }

    fn print_expression(&self, expression: &Expression) -> Result<()> {
        use Expression::*;
        match expression {
            Closure(closure) => self.print_closure(closure),
            Identifier(identifier) => self.print_identifier(identifier),
            Infix(infix) => self.print_infix(infix),
            LiteralInt(literal) => self.lnwrite(format!("{}", literal.value)),
            PostfixCall(call) => self.print_postfix_call(call, 0).map(|_| ()),
            PostfixProperty(property) => self.print_postfix_property(property, 0).map(|_| ()),
        }
    }

    fn print_closure(&self, closure: &Closure) -> Result<()> {
        self.lnwrite("Closure(")?;
        if !closure.arguments.is_empty() {
            self.write("\n")?;
            self.indented(|this| {
                for argument in closure.arguments.iter() {
                    this.iwrite(format!("{}: ", argument.name))?;
                    this.write_type(&argument.typ, true)?;
                    this.write(",\n")?;
                }
                Ok(())
            })?;
        }
        self.iwrite("): ")?;
        let retrn = match &closure.typ {
            Type::Func(func) => func.retrn.borrow(),
            other @ _ => unreachable!("Closure node has non-Func type: {:?}", other),
        };
        self.write_type(&retrn, false)?;
        self.write(" ")?;
        match &*closure.body {
            ClosureBody::Block(block) => self.print_block(block, false),
            ClosureBody::Expression(expression) => self.indented(|this| {
                this.write("(")?;
                this.print_expression(expression)?;
                this.write(")")
            }),
        }
    }

    fn print_identifier(&self, identifier: &Identifier) -> Result<()> {
        self.lnwrite(format!("Identifier({}): ", identifier.name.name))?;
        self.write_type(&identifier.typ, false)
    }

    fn print_infix(&self, infix: &Infix) -> Result<()> {
        self.lnwrite("Infix(")?;
        self.indented(|this1| {
            this1.lnwrite("lhs:")?;
            this1.indented(|this2| this2.print_expression(&infix.lhs))?;
            this1.lnwrite(format!("op: {}", infix.op.to_string()))?;
            this1.lnwrite(format!("rhs:"))?;
            this1.indented(|this2| this2.print_expression(&infix.rhs))
        })?;
        self.write(")")
    }

    fn is_postfix(&self, expression: &Expression) -> bool {
        use Expression::*;
        match expression {
            PostfixProperty(_) => true,
            _ => false,
        }
    }

    fn print_postfix_call(&self, call: &PostfixCall, current: u8) -> Result<u8> {
        let max = self.print_postfix_target(&*call.target, current)?;
        // Have a +1 so that we always indent at least one step.
        self.indented_steps(max - current + 1, |this1| {
            this1.lnwrite(format!("Call("))?;
            if !call.arguments.is_empty() {
                for argument in call.arguments.iter() {
                    this1.indented(|this2| {
                        this2.print_expression(argument)?;
                        this2.write(",")
                    })?;
                }
            }
            this1.write("): ")?;
            this1.write_type(&call.typ, false)
        })?;
        Ok(max)
    }

    fn print_postfix_property(&self, property: &PostfixProperty, current: u8) -> Result<u8> {
        let max = self.print_postfix_target(&*property.target, current)?;
        self.indented_steps(max - current + 1, |this| {
            this.lnwrite(format!("Property({}): ", property.property.name))?;
            this.write_type(&property.typ, false)
        })?;
        Ok(max)
    }

    fn print_postfix_target(&self, target: &Expression, current: u8) -> Result<u8> {
        match target {
            // Links in the chain
            Expression::PostfixCall(target) => self.print_postfix_call(target, current + 1),
            Expression::PostfixProperty(target) => self.print_postfix_property(target, current + 1),
            // Tail of the chain
            other @ _ => {
                self.print_expression(other)?;
                Ok(current)
            }
        }
    }

    /// Write a string.
    fn write<S: AsRef<str>>(&self, string: S) -> Result<()> {
        self.write_output(string.as_ref().as_bytes())
    }

    /// Write indentation and then a string.
    fn iwrite<S: AsRef<str>>(&self, string: S) -> Result<()> {
        let indented = " ".repeat(self.indent.get() as usize);
        self.write_output(indented.as_bytes())
            .and_then(|_| self.write_output(string.as_ref().as_bytes()))
            .map(|_| ())
    }

    /// Write a newline, indentation, and then a string.
    fn lnwrite<S: AsRef<str>>(&self, string: S) -> Result<()> {
        let indented = " ".repeat(self.indent.get() as usize);
        self.write_output("\n".as_bytes())
            .and_then(|_| self.write_output(indented.as_bytes()))
            .and_then(|_| self.write_output(string.as_ref().as_bytes()))
    }

    fn write_output(&self, bytes: &[u8]) -> Result<()> {
        let mut output = self.output.borrow_mut();
        output.write(bytes).map(|_| ())
    }

    fn write_pointer<P>(&self, ptr: *const P) -> Result<()> {
        if !self.print_pointers {
            return Ok(());
        }
        self.write(format!(" @ {:p}", ptr))
    }

    // fn write_output(&self, bytes: &[u8]) -> Result<()> {
    //     let result = if let Some(buffer) = &mut *self.buffer.borrow_mut() {
    //         buffer.write(bytes)
    //     } else {
    //         let mut output = self.output.borrow_mut();
    //         output.write(bytes)
    //     };
    //     result.map(|_| ())
    // }
    //
    // fn buffer<F: FnOnce(&Self) -> Result<()>>(&self, inner: F) -> Result<Vec<u8>> {
    //     {
    //         let mut buffer = self.buffer.borrow_mut();
    //         *buffer = Some(vec![]);
    //     }
    //     inner(self)?;
    //     let buffer = self.buffer.replace(None);
    //     Ok(buffer.unwrap())
    // }
}
<|endoftext|>


<commit_before>use crate::grammar::parsers::testing::setup;
use crate::grammar::model::Fragment;
use nom::IResult;
use std::mem::discriminant;
use std::ptr::eq as ptr_eq;
use std::fmt::Debug;
use nom::Err;
use nom::bytes::complete::take_while1;
use nom::character::is_alphabetic;
use codespan::Files;

fn test_frag_str<'a, 'b, F1, F2>(
    s: &'static str,
    f1: F1,
    f2: F2)
where
    F1: Fn(Fragment<'a>) -> IResult<Fragment<'a>, Fragment<'a>>,
    F2: Fn(&'b str) -> IResult<&'b str, &'b str>,
{
    let mut f: Files<String> = Files::new();
    let h = f.add("test", s.to_owned());
    let frag: Fragment = Fragment::new(&f, h);
    let p1 = f1(frag);
    let p2 = f2(s);
    match (p1, p2) {
        (Ok((rem1, out1)), Ok((rem2, out2))) => {
            assert_eq!(rem1.source(), rem2);
            assert_eq!(out1.source(), out2);
            assert_eq!(out1.len(), out2.len());
            assert_eq!(out1.source().len(), out1.len());
            assert!(ptr_eq(rem1.files(), frag.files()));
            assert_eq!(rem1.len(), rem2.len());
            assert_eq!(rem1.get_handle(), frag.get_handle());
        },
        (   Result::Err(Err::Error((fr, ek1))),
            Result::Err(Err::Error((str, ek2)))
        ) => {
            let fr: Fragment = fr;
            assert_eq!(ek1, ek2);
            assert_eq!(fr.get_handle(), frag.get_handle());
            assert!(ptr_eq(fr.files(), frag.files()));
            assert_eq!(fr.len(), str.len());
            assert_eq!(fr.source().len(), fr.len());
            assert_eq!(fr.source(), str);
        },
        (fr, s) =>
            panic!("Parser output mismatch!\nFrag parser:\n{:?}\nStr parser:\n{:?}", fr, s)
    }
}

#[test]
fn test_take_while1() {
    let is_alpha = move |c: char| c.is_alphabetic();
    let ff = take_while1(is_alpha);
    let fs = take_while1( is_alpha);
    test_frag_str("", ff, fs);
}<commit_msg>borrowch<commit_after>use crate::grammar::parsers::testing::setup;
use crate::grammar::model::Fragment;
use nom::IResult;
use std::mem::discriminant;
use std::ptr::eq as ptr_eq;
use std::fmt::Debug;
use nom::Err;
use nom::bytes::complete::take_while1;
use nom::character::is_alphabetic;
use codespan::Files;

fn test_frag_str<'a, 'b, F1, F2>(
    s: &'static str,
    f1: F1,
    f2: F2)
where
    F1: for<'d> Fn(Fragment<'d>) -> IResult<Fragment<'d>, Fragment<'d>>,
    F2: Fn(&'b str) -> IResult<&'b str, &'b str>,
{
    let mut f: Files<String> = Files::new();
    let h = f.add("test", s.to_owned());
    let frag: Fragment = Fragment::new(&f, h);
    let p1 = f1(frag);
    let p2 = f2(s);
    match (p1, p2) {
        (Ok((rem1, out1)), Ok((rem2, out2))) => {
            assert_eq!(rem1.source(), rem2);
            assert_eq!(out1.source(), out2);
            assert_eq!(out1.len(), out2.len());
            assert_eq!(out1.source().len(), out1.len());
            assert!(ptr_eq(rem1.files(), frag.files()));
            assert_eq!(rem1.len(), rem2.len());
            assert_eq!(rem1.get_handle(), frag.get_handle());
        },
        (   Result::Err(Err::Error((fr, ek1))),
            Result::Err(Err::Error((str, ek2)))
        ) => {
            let fr: Fragment = fr;
            assert_eq!(ek1, ek2);
            assert_eq!(fr.get_handle(), frag.get_handle());
            assert!(ptr_eq(fr.files(), frag.files()));
            assert_eq!(fr.len(), str.len());
            assert_eq!(fr.source().len(), fr.len());
            assert_eq!(fr.source(), str);
        },
        (fr, s) =>
            panic!("Parser output mismatch!\nFrag parser:\n{:?}\nStr parser:\n{:?}", fr, s)
    }
}

#[test]
fn test_take_while1() {
    let is_alpha = move |c: char| c.is_alphabetic();
    let ff = |input: Fragment| take_while1(is_alpha)(input);
    let fs = take_while1( is_alpha);
    test_frag_str("", ff, fs);
}<|endoftext|>


<commit_before>#![crate_name = "docopt_macros"]
#![crate_type = "dylib"]

#![feature(plugin_registrar, quote, rustc_private)]

//! This crate defines the `docopt!` macro. It is documented in the
//! documentation of the `docopt` crate.

extern crate syntax;
extern crate rustc_plugin;
extern crate docopt;

use std::borrow::Borrow;
use std::collections::HashMap;

use rustc_plugin::Registry;
use syntax::{ast, codemap, tokenstream};
use syntax::errors::DiagnosticBuilder;
use syntax::ext::base::{ExtCtxt, MacResult, MacEager, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::parse::common::SeqSep;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::print::pprust;
use syntax::ptr::P;
use syntax::util::small_vector::SmallVector;

use docopt::{Docopt, ArgvMap};
use docopt::parse::{Options, Atom, Positional, Zero, One};

type PResult<'a, T> = Result<T, DiagnosticBuilder<'a>>;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("docopt", expand);
}

fn expand(cx: &mut ExtCtxt, span: codemap::Span, tts: &[tokenstream::TokenTree])
         -> Box<MacResult+'static> {
    let parsed = match MacParser::new(cx, tts).parse() {
        Ok(parsed) => parsed,
        Err(_) => return DummyResult::any(span),
    };
    parsed.items(cx)
}

/// Parsed corresponds to the result of parsing a `docopt` macro call.
/// It can be used to write a corresponding struct.
struct Parsed {
    struct_info: StructInfo,
    doc: Docopt,
    /// Overrided type annotations for struct members. May be empty.
    /// When a type annotation for an atom doesn't exist, then one is
    /// inferred automatically. It is one of: `bool`, `usize`, `String` or
    /// `Vec<String>`.
    types: HashMap<Atom, P<ast::Ty>>,
}

impl Parsed {
    /// Returns a macro result suitable for expansion.
    /// Contains two items: one for the struct and one for the struct impls.
    fn items(&self, cx: &ExtCtxt) -> Box<MacResult+'static> {
        let mut its = vec!();
        its.push(self.struct_decl(cx));

        let struct_name = self.struct_info.name;
        let full_doc = &*self.doc.parser().full_doc;
        its.push(quote_item!(cx,
            impl $struct_name {
                #[allow(dead_code)]
                fn docopt() -> docopt::Docopt {
                    // The unwrap is justified here because this code
                    // gen only happens if the Docopt usage string is valid.
                    docopt::Docopt::new($full_doc).unwrap()
                }
            }
        ).unwrap());
        MacEager::items(SmallVector::many(its))
    }

    /// Returns an item for the struct definition.
    fn struct_decl(&self, cx: &ExtCtxt) -> P<ast::Item> {
        let name = self.struct_info.name.clone();
        let vis = if self.struct_info.public { ast::Visibility::Public }
                  else { ast::Visibility::Inherited };
        let def = ast::VariantData::Struct(
            self.struct_fields(cx), ast::DUMMY_NODE_ID);

        let mut traits = vec!["RustcDecodable".to_string()];
        traits.extend(self.struct_info.deriving.iter().cloned());
        let attrs = vec![attribute(cx, "allow", vec!["non_snake_case"]),
                         attribute(cx, "derive", traits)];

        let st = cx.item_struct(codemap::DUMMY_SP, name.clone(), def);
        cx.item(codemap::DUMMY_SP, name, attrs, st.node.clone()).map(|mut it| {
            it.vis = vis;
            it
        })
    }

    /// Returns a list of fields for the struct definition.
    /// Handles type annotations.
    fn struct_fields(&self, cx: &ExtCtxt) -> Vec<ast::StructField> {
        let mut fields: Vec<ast::StructField> = vec!();
        for (atom, opts) in self.doc.parser().descs.iter() {
            let name = ArgvMap::key_to_struct_field(&*atom.to_string());
            let ty = match self.types.get(atom) {
                None => self.pat_type(cx, atom, opts),
                Some(ty) => ty.clone(),
            };
            fields.push(self.mk_struct_field(&*name, ty));
        }
        fields
    }

    /// Returns an inferred type for a usage pattern.
    /// This is only invoked when a type annotation is not present.
    fn pat_type(&self, cx: &ExtCtxt, atom: &Atom, opts: &Options) -> P<ast::Ty> {
        let sp = codemap::DUMMY_SP;
        match (opts.repeats, &opts.arg) {
            (false, &Zero) => {
                match atom {
                    &Positional(_) => cx.ty_ident(sp, ident("String")),
                    _ => cx.ty_ident(sp, ident("bool")),
                }
            }
            (true, &Zero) => {
                match atom {
                    &Positional(_) => ty_vec_string(cx),
                    _ => cx.ty_ident(sp, ident("usize")),
                }
            }
            (false, &One(_)) => cx.ty_ident(sp, ident("String")),
            (true, &One(_)) => ty_vec_string(cx),
        }
    }

    /// Creates a struct field from a member name and type.
    fn mk_struct_field(&self, name: &str, ty: P<ast::Ty>) -> ast::StructField {
        ast::StructField {
            span: codemap::DUMMY_SP,
            ident: Some(ident(name)),
            vis: ast::Visibility::Public,
            id: ast::DUMMY_NODE_ID,
            ty: ty,
            attrs: vec!(),
        }
    }
}

/// State for parsing a `docopt` macro invocation.
struct MacParser<'a, 'b:'a> {
    cx: &'a mut ExtCtxt<'b>,
    p: Parser<'b>,
}

impl<'a, 'b> MacParser<'a, 'b> {
    fn new(cx: &'a mut ExtCtxt<'b>, tts: &[tokenstream::TokenTree]) -> MacParser<'a, 'b> {
        let p = cx.new_parser_from_tts(tts);
        MacParser { cx: cx, p: p }
    }

    /// Main entry point for parsing arguments to `docopt` macro.
    /// First looks for an identifier for the struct name.
    /// Second, a string containing the docopt usage patterns.
    /// Third, an optional list of type annotations.
    fn parse(&mut self) -> PResult<'b, Parsed> {
        if self.p.token == token::Eof {
            let err = self.cx.struct_span_err(
                self.cx.call_site(), "macro expects arguments");
            return Err(err);
        }
        let struct_info = try!(self.parse_struct_info());
        let docstr = try!(self.parse_str());

        let mut types = HashMap::new();
        if !self.p.check(&token::Eof) {
            let sep = SeqSep {
                sep: Some(token::Comma),
                trailing_sep_allowed: true,
            };
            types = self.p.parse_seq_to_before_end(
                &token::Eof, sep, |p| MacParser::parse_type_annotation(p)
            ).into_iter()
             .map(|(ident, ty)| {
                  let field_name = ident.to_string();
                  let key = ArgvMap::struct_field_to_key(&*field_name);
                  (Atom::new(&*key), ty)
              })
             .collect::<HashMap<Atom, P<ast::Ty>>>();
            try!(self.p.expect(&token::Eof));
        }

        // This config does not matter because we're only asking for the
        // usage patterns in the Docopt string. The configuration does not
        // affect the retrieval of usage patterns.
        let doc = match Docopt::new(docstr) {
            Ok(doc) => doc,
            Err(err) => {
                let err = self.cx.struct_span_err(
                    self.cx.call_site(),
                    &format!("Invalid Docopt usage: {}", err));
                return Err(err);
            }
        };
        Ok(Parsed {
            struct_info: struct_info,
            doc: doc,
            types: types,
        })
    }

    /// Parses a single string literal. On failure, an error is logged and
    /// unit is returned.
    fn parse_str(&mut self) -> PResult<'b, String> {
        fn lit_is_str(lit: &ast::Lit) -> bool {
            match lit.node {
                ast::LitKind::Str(_, _) => true,
                _ => false,
            }
        }
        fn lit_to_string(lit: &ast::Lit) -> String {
            match lit.node {
                ast::LitKind::Str(ref s, _) => s.to_string(),
                _ => panic!("BUG: expected string literal"),
            }
        }
        let exp = self.cx.expander().fold_expr(self.p.parse_expr().unwrap());
        let s = match exp.node {
            ast::ExprKind::Lit(ref lit) if lit_is_str(&**lit) => {
                lit_to_string(&**lit)
            }
            _ => {
                let err = format!("Expected string literal but got {}",
                                  pprust::expr_to_string(&*exp));
                let err = self.cx.struct_span_err(exp.span, &*err);
                return Err(err);
            }
        };
        self.p.bump();
        Ok(s)
    }

    /// Parses a type annotation in a `docopt` invocation of the form
    /// `ident: Ty`.
    /// Note that this is a static method as it is used as a HOF.
    fn parse_type_annotation(p: &mut Parser<'b>)
                             -> PResult<'b, (ast::Ident, P<ast::Ty>)> {
        let ident = try!(p.parse_ident());
        try!(p.expect(&token::Colon));
        let ty = p.parse_ty().unwrap();
        Ok((ident, ty))
    }

    /// Parses struct information, like visibility, name and deriving.
    fn parse_struct_info(&mut self) -> PResult<'b, StructInfo> {
        let public = self.p.eat_keyword(token::keywords::Pub);
        let mut info = StructInfo {
            name: try!(self.p.parse_ident()),
            public: public,
            deriving: vec![],
        };
        if self.p.eat(&token::Comma) { return Ok(info); }
        let deriving = try!(self.p.parse_ident());
        if deriving.name.as_str() != "derive" {
            let err = format!("Expected 'derive' keyword but got '{}'",
                              deriving);
            let err = self.cx.struct_span_err(self.cx.call_site(), &*err);
            return Err(err);
        }
        while !self.p.eat(&token::Comma) {
            info.deriving.push(
                try!(self.p.parse_ident()).name.to_string());
        }
        Ok(info)
    }
}

struct StructInfo {
    name: ast::Ident,
    public: bool,
    deriving: Vec<String>,
}

// Convenience functions for building intermediate values.

fn ident(s: &str) -> ast::Ident {
    ast::Ident::with_empty_ctxt(token::intern(s))
}

fn attribute<S, T>(cx: &ExtCtxt, name: S, items: Vec<T>) -> ast::Attribute
            where S: Borrow<str>, T: Borrow<str> {
    let sp = codemap::DUMMY_SP;
    let its = items.into_iter().map(|s| meta_item(cx, s.borrow())).collect();
    let mi = cx.meta_list(sp, intern(name.borrow()), its);
    cx.attribute(sp, mi)
}

fn meta_item(cx: &ExtCtxt, s: &str) -> codemap::Spanned<ast::NestedMetaItemKind> {
    codemap::Spanned {
        node: ast::NestedMetaItemKind::MetaItem(cx.meta_word(codemap::DUMMY_SP, intern(s))),
        span: cx.call_site(),
    }
}

fn intern(s: &str) -> token::InternedString {
    token::intern_and_get_ident(s)
}

fn ty_vec_string(cx: &ExtCtxt) -> P<ast::Ty> {
    let sp = codemap::DUMMY_SP;
    let tystr = ast::AngleBracketedParameterData {
        lifetimes: vec![],
        types: P::from_vec(vec![cx.ty_ident(sp, ident("String"))]),
        bindings: P::new(),
    };
    cx.ty_path(ast::Path {
        span: sp,
        global: false,
        segments: vec![ast::PathSegment {
            identifier: ident("Vec"),
            parameters: ast::PathParameters::AngleBracketed(tystr),
        }]
    })
}
<commit_msg>update for rust nightly<commit_after>#![crate_name = "docopt_macros"]
#![crate_type = "dylib"]

#![feature(plugin_registrar, quote, rustc_private)]

//! This crate defines the `docopt!` macro. It is documented in the
//! documentation of the `docopt` crate.

extern crate syntax;
extern crate rustc_plugin;
extern crate docopt;

use std::borrow::Borrow;
use std::collections::HashMap;

use rustc_plugin::Registry;
use syntax::{ast, codemap, tokenstream};
use syntax::errors::DiagnosticBuilder;
use syntax::ext::base::{ExtCtxt, MacResult, MacEager, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::parse::common::SeqSep;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::print::pprust;
use syntax::ptr::P;
use syntax::symbol;
use syntax::util::small_vector::SmallVector;

use docopt::{Docopt, ArgvMap};
use docopt::parse::{Options, Atom, Positional, Zero, One};

type PResult<'a, T> = Result<T, DiagnosticBuilder<'a>>;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("docopt", expand);
}

fn expand(cx: &mut ExtCtxt, span: codemap::Span, tts: &[tokenstream::TokenTree])
         -> Box<MacResult+'static> {
    let parsed = match MacParser::new(cx, tts).parse() {
        Ok(parsed) => parsed,
        Err(_) => return DummyResult::any(span),
    };
    parsed.items(cx)
}

/// Parsed corresponds to the result of parsing a `docopt` macro call.
/// It can be used to write a corresponding struct.
struct Parsed {
    struct_info: StructInfo,
    doc: Docopt,
    /// Overrided type annotations for struct members. May be empty.
    /// When a type annotation for an atom doesn't exist, then one is
    /// inferred automatically. It is one of: `bool`, `usize`, `String` or
    /// `Vec<String>`.
    types: HashMap<Atom, P<ast::Ty>>,
}

impl Parsed {
    /// Returns a macro result suitable for expansion.
    /// Contains two items: one for the struct and one for the struct impls.
    fn items(&self, cx: &ExtCtxt) -> Box<MacResult+'static> {
        let mut its = vec!();
        its.push(self.struct_decl(cx));

        let struct_name = self.struct_info.name;
        let full_doc = &*self.doc.parser().full_doc;
        its.push(quote_item!(cx,
            impl $struct_name {
                #[allow(dead_code)]
                fn docopt() -> docopt::Docopt {
                    // The unwrap is justified here because this code
                    // gen only happens if the Docopt usage string is valid.
                    docopt::Docopt::new($full_doc).unwrap()
                }
            }
        ).unwrap());
        MacEager::items(SmallVector::many(its))
    }

    /// Returns an item for the struct definition.
    fn struct_decl(&self, cx: &ExtCtxt) -> P<ast::Item> {
        let name = self.struct_info.name.clone();
        let vis = if self.struct_info.public { ast::Visibility::Public }
                  else { ast::Visibility::Inherited };
        let def = ast::VariantData::Struct(
            self.struct_fields(cx), ast::DUMMY_NODE_ID);

        let mut traits = vec!["RustcDecodable".to_string()];
        traits.extend(self.struct_info.deriving.iter().cloned());
        let attrs = vec![attribute(cx, "allow", vec!["non_snake_case"]),
                         attribute(cx, "derive", traits)];

        let st = cx.item_struct(codemap::DUMMY_SP, name.clone(), def);
        cx.item(codemap::DUMMY_SP, name, attrs, st.node.clone()).map(|mut it| {
            it.vis = vis;
            it
        })
    }

    /// Returns a list of fields for the struct definition.
    /// Handles type annotations.
    fn struct_fields(&self, cx: &ExtCtxt) -> Vec<ast::StructField> {
        let mut fields: Vec<ast::StructField> = vec!();
        for (atom, opts) in self.doc.parser().descs.iter() {
            let name = ArgvMap::key_to_struct_field(&*atom.to_string());
            let ty = match self.types.get(atom) {
                None => self.pat_type(cx, atom, opts),
                Some(ty) => ty.clone(),
            };
            fields.push(self.mk_struct_field(&*name, ty));
        }
        fields
    }

    /// Returns an inferred type for a usage pattern.
    /// This is only invoked when a type annotation is not present.
    fn pat_type(&self, cx: &ExtCtxt, atom: &Atom, opts: &Options) -> P<ast::Ty> {
        let sp = codemap::DUMMY_SP;
        match (opts.repeats, &opts.arg) {
            (false, &Zero) => {
                match atom {
                    &Positional(_) => cx.ty_ident(sp, ident("String")),
                    _ => cx.ty_ident(sp, ident("bool")),
                }
            }
            (true, &Zero) => {
                match atom {
                    &Positional(_) => ty_vec_string(cx),
                    _ => cx.ty_ident(sp, ident("usize")),
                }
            }
            (false, &One(_)) => cx.ty_ident(sp, ident("String")),
            (true, &One(_)) => ty_vec_string(cx),
        }
    }

    /// Creates a struct field from a member name and type.
    fn mk_struct_field(&self, name: &str, ty: P<ast::Ty>) -> ast::StructField {
        ast::StructField {
            span: codemap::DUMMY_SP,
            ident: Some(ident(name)),
            vis: ast::Visibility::Public,
            id: ast::DUMMY_NODE_ID,
            ty: ty,
            attrs: vec!(),
        }
    }
}

/// State for parsing a `docopt` macro invocation.
struct MacParser<'a, 'b:'a> {
    cx: &'a mut ExtCtxt<'b>,
    p: Parser<'b>,
}

impl<'a, 'b> MacParser<'a, 'b> {
    fn new(cx: &'a mut ExtCtxt<'b>, tts: &[tokenstream::TokenTree]) -> MacParser<'a, 'b> {
        let p = cx.new_parser_from_tts(tts);
        MacParser { cx: cx, p: p }
    }

    /// Main entry point for parsing arguments to `docopt` macro.
    /// First looks for an identifier for the struct name.
    /// Second, a string containing the docopt usage patterns.
    /// Third, an optional list of type annotations.
    fn parse(&mut self) -> PResult<'b, Parsed> {
        if self.p.token == token::Eof {
            let err = self.cx.struct_span_err(
                self.cx.call_site(), "macro expects arguments");
            return Err(err);
        }
        let struct_info = try!(self.parse_struct_info());
        let docstr = try!(self.parse_str());

        let mut types = HashMap::new();
        if !self.p.check(&token::Eof) {
            let sep = SeqSep {
                sep: Some(token::Comma),
                trailing_sep_allowed: true,
            };
            types = self.p.parse_seq_to_before_end(
                &token::Eof, sep, |p| MacParser::parse_type_annotation(p)
            ).into_iter()
             .map(|(ident, ty)| {
                  let field_name = ident.to_string();
                  let key = ArgvMap::struct_field_to_key(&*field_name);
                  (Atom::new(&*key), ty)
              })
             .collect::<HashMap<Atom, P<ast::Ty>>>();
            try!(self.p.expect(&token::Eof));
        }

        // This config does not matter because we're only asking for the
        // usage patterns in the Docopt string. The configuration does not
        // affect the retrieval of usage patterns.
        let doc = match Docopt::new(docstr) {
            Ok(doc) => doc,
            Err(err) => {
                let err = self.cx.struct_span_err(
                    self.cx.call_site(),
                    &format!("Invalid Docopt usage: {}", err));
                return Err(err);
            }
        };
        Ok(Parsed {
            struct_info: struct_info,
            doc: doc,
            types: types,
        })
    }

    /// Parses a single string literal. On failure, an error is logged and
    /// unit is returned.
    fn parse_str(&mut self) -> PResult<'b, String> {
        fn lit_is_str(lit: &ast::Lit) -> bool {
            match lit.node {
                ast::LitKind::Str(_, _) => true,
                _ => false,
            }
        }
        fn lit_to_string(lit: &ast::Lit) -> String {
            match lit.node {
                ast::LitKind::Str(ref s, _) => s.to_string(),
                _ => panic!("BUG: expected string literal"),
            }
        }
        let exp = self.cx.expander().fold_expr(self.p.parse_expr().unwrap());
        let s = match exp.node {
            ast::ExprKind::Lit(ref lit) if lit_is_str(&**lit) => {
                lit_to_string(&**lit)
            }
            _ => {
                let err = format!("Expected string literal but got {}",
                                  pprust::expr_to_string(&*exp));
                let err = self.cx.struct_span_err(exp.span, &*err);
                return Err(err);
            }
        };
        self.p.bump();
        Ok(s)
    }

    /// Parses a type annotation in a `docopt` invocation of the form
    /// `ident: Ty`.
    /// Note that this is a static method as it is used as a HOF.
    fn parse_type_annotation(p: &mut Parser<'b>)
                             -> PResult<'b, (ast::Ident, P<ast::Ty>)> {
        let ident = try!(p.parse_ident());
        try!(p.expect(&token::Colon));
        let ty = p.parse_ty().unwrap();
        Ok((ident, ty))
    }

    /// Parses struct information, like visibility, name and deriving.
    fn parse_struct_info(&mut self) -> PResult<'b, StructInfo> {
        let public = self.p.eat_keyword(symbol::keywords::Pub);
        let mut info = StructInfo {
            name: try!(self.p.parse_ident()),
            public: public,
            deriving: vec![],
        };
        if self.p.eat(&token::Comma) { return Ok(info); }
        let deriving = try!(self.p.parse_ident());
        if *deriving.name.as_str() != *"derive" {
            let err = format!("Expected 'derive' keyword but got '{}'",
                              deriving);
            let err = self.cx.struct_span_err(self.cx.call_site(), &*err);
            return Err(err);
        }
        while !self.p.eat(&token::Comma) {
            info.deriving.push(
                try!(self.p.parse_ident()).name.to_string());
        }
        Ok(info)
    }
}

struct StructInfo {
    name: ast::Ident,
    public: bool,
    deriving: Vec<String>,
}

// Convenience functions for building intermediate values.

fn ident(s: &str) -> ast::Ident {
    ast::Ident::with_empty_ctxt(symbol::Symbol::intern(s))
}

fn attribute<S, T>(cx: &ExtCtxt, name: S, items: Vec<T>) -> ast::Attribute
            where S: Borrow<str>, T: Borrow<str> {
    let sp = codemap::DUMMY_SP;
    let its = items.into_iter().map(|s| meta_item(cx, s.borrow())).collect();
    let mi = cx.meta_list(sp, intern(name.borrow()), its);
    cx.attribute(sp, mi)
}

fn meta_item(cx: &ExtCtxt, s: &str) -> codemap::Spanned<ast::NestedMetaItemKind> {
    codemap::Spanned {
        node: ast::NestedMetaItemKind::MetaItem(cx.meta_word(codemap::DUMMY_SP, intern(s))),
        span: cx.call_site(),
    }
}

fn intern(s: &str) -> symbol::Symbol {
    symbol::Symbol::intern(s)
}

fn ty_vec_string(cx: &ExtCtxt) -> P<ast::Ty> {
    let sp = codemap::DUMMY_SP;
    let tystr = ast::AngleBracketedParameterData {
        lifetimes: vec![],
        types: P::from_vec(vec![cx.ty_ident(sp, ident("String"))]),
        bindings: P::new(),
    };
    cx.ty_path(ast::Path {
        span: sp,
        global: false,
        segments: vec![ast::PathSegment {
            identifier: ident("Vec"),
            parameters: ast::PathParameters::AngleBracketed(tystr),
        }]
    })
}
<|endoftext|>