<commit_before>//! Network related functionality for `KafkaClient`.
//!
//! This module is crate private and not exposed to the public except
//! through re-exports of individual items from within
//! `kafka::client`.

use std::collections::HashMap;
use std::fmt;
use std::io::{Read, Write};
use std::mem;
use std::net::{TcpStream, Shutdown};
use std::time::{Instant, Duration};

#[cfg(feature = "security")]
use openssl::ssl::{SslContext, SslStream};

use error::Result;

// --------------------------------------------------------------------

/// Security relevant configuration options for `KafkaClient`.
// This will be expanded in the future. See #51.
#[cfg(feature = "security")]
#[derive(Debug)]
pub struct SecurityConfig(SslContext);

#[cfg(feature = "security")]
impl SecurityConfig {
    /// In the future this will also support a kerbos via #51.
    pub fn new(ssl: SslContext) -> SecurityConfig {
        SecurityConfig(ssl)
    }
}

// --------------------------------------------------------------------

struct Pooled<T> {
    last_checkout: Instant,
    item: T,
}

impl<T> Pooled<T> {
    fn new(last_checkout: Instant, item: T) -> Self {
        Pooled { last_checkout: last_checkout, item: item }
    }
}

impl<T: fmt::Debug> fmt::Debug for Pooled<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Pooled {{ last_checkout: {:?}, item: {:?} }}",
               self.last_checkout, self.item)
    }
}

#[derive(Debug)]
pub struct Config {
    rw_timeout: Option<Duration>,
    idle_timeout: Duration,
    #[cfg(feature = "security")]
    security_config: Option<SecurityConfig>,
}

impl Config {
    #[cfg(not(feature = "security"))]
    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
        KafkaConnection::new(id, host, self.rw_timeout)
            .map(|c| {
                debug!("Established: {:?}", c);
                c
            })
    }

    #[cfg(feature = "security")]
    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
        KafkaConnection::new(id,
                             host,
                             self.rw_timeout,
                             self.security_config.as_ref().map(|c| &c.0))
            .map(|c| {
                debug!("Established: {:?}", c);
                c
            })
    }
}

#[derive(Debug)]
struct State {
    num_conns: u32,
}

impl State {
    fn new() -> State {
        State { num_conns: 0 }
    }

    fn next_conn_id(&mut self) -> u32 {
        let c = self.num_conns;
        self.num_conns = self.num_conns.wrapping_add(1);
        c
    }
}

#[derive(Debug)]
pub struct Connections {
    conns: HashMap<String, Pooled<KafkaConnection>>,
    state: State,
    config: Config,
}

impl Connections {
    #[cfg(not(feature = "security"))]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
        Connections {
            conns: HashMap::new(),
            state: State::new(),
            config: Config {
                rw_timeout: rw_timeout,
                idle_timeout: idle_timeout,
            },
        }
    }

    #[cfg(feature = "security")]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
        Self::new_with_security(rw_timeout, idle_timeout, None)
    }

    #[cfg(feature = "security")]
    pub fn new_with_security(rw_timeout: Option<Duration>,
                             idle_timeout: Duration,
                             security: Option<SecurityConfig>)
                             -> Connections
    {
        Connections {
            conns: HashMap::new(),
            state: State::new(),
            config: Config {
                rw_timeout: rw_timeout,
                idle_timeout: idle_timeout,
                security_config: security,
            },
        }
    }

    pub fn set_idle_timeout(&mut self, idle_timeout: Duration) {
        self.config.idle_timeout = idle_timeout;
    }

    pub fn idle_timeout(&self) -> Duration {
        self.config.idle_timeout
    }

    pub fn get_conn<'a>(&'a mut self, host: &str, now: Instant)
                        -> Result<&'a mut KafkaConnection>
    {
        if let Some(conn) = self.conns.get_mut(host) {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout reached: {:?}", conn.item);
                let new_conn = try!(self.config.new_conn(self.state.next_conn_id(), host));
                let _ = conn.item.shutdown();
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            // ~ decouple the lifetimes to make the borrowck happy;
            // this is safe since we're immediatelly returning the
            // reference and the rest of the code in this method is
            // not affected
            return Ok(unsafe { mem::transmute(kconn) });
        }
        let cid = self.state.next_conn_id();
        self.conns.insert(host.to_owned(),
                          Pooled::new(now, try!(self.config.new_conn(cid, host))));
        Ok(&mut self.conns.get_mut(host).unwrap().item)
    }

    pub fn get_conn_any(&mut self, now: Instant) -> Option<&mut KafkaConnection> {
        for (host, conn) in &mut self.conns {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout reached: {:?}", conn.item);
                let new_conn_id = self.state.next_conn_id();
                let new_conn = match self.config.new_conn(new_conn_id, host.as_str()) {
                    Ok(new_conn) => {
                        let _ = conn.item.shutdown();
                        new_conn
                    }
                    Err(e) => {
                        warn!("Failed to establish connection to {}: {:?}", host, e);
                        continue;
                    }
                };
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            return Some(kconn);
        }
        None
    }
}

// --------------------------------------------------------------------

trait IsSecured {
    fn is_secured(&self) -> bool;
}

#[cfg(not(feature = "security"))]
type KafkaStream = TcpStream;

#[cfg(not(feature = "security"))]
impl IsSecured for KafkaStream {
    fn is_secured(&self) -> bool {
        false
    }
}

#[cfg(feature = "security")]
use self::openssled::KafkaStream;

#[cfg(feature = "security")]
mod openssled {
    use std::io::{self, Read, Write};
    use std::net::{TcpStream, Shutdown};
    use std::time::Duration;

    use openssl::ssl::SslStream;

    use super::IsSecured;

    pub enum KafkaStream {
        Plain(TcpStream),
        Ssl(SslStream<TcpStream>)
    }

    impl IsSecured for KafkaStream {
        fn is_secured(&self) -> bool {
            match *self {
                KafkaStream::Ssl(_) => true,
                _ => false,
            }
        }
    }

    impl KafkaStream {
        fn get_ref(&self) -> &TcpStream {
            match *self {
                KafkaStream::Plain(ref s) => s,
                KafkaStream::Ssl(ref s) => s.get_ref()
            }
        }

        pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_read_timeout(dur)
        }

        pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_write_timeout(dur)
        }

        pub fn shutdown(&mut self, how: Shutdown) -> io::Result<()> {
            self.get_ref().shutdown(how)
        }
    }

    impl Read for KafkaStream {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.read(buf),
                KafkaStream::Ssl(ref mut s) => s.read(buf),
            }
        }
    }

    impl Write for KafkaStream {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.write(buf),
                KafkaStream::Ssl(ref mut s) => s.write(buf),
            }
        }
        fn flush(&mut self) -> io::Result<()> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.flush(),
                KafkaStream::Ssl(ref mut s) => s.flush(),
            }
        }
    }
}

/// A TCP stream to a remote Kafka broker.
pub struct KafkaConnection {
    // a surrogate identifier to distinguish between
    // connections to the same host in debug messages
    id: u32,
    // "host:port"
    host: String,
    // the (wrapped) tcp stream
    stream: KafkaStream,
}

impl fmt::Debug for KafkaConnection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "KafkaConnection {{ id: {}, secured: {}, host: \"{}\" }}",
               self.id, self.stream.is_secured(), self.host)
    }
}

impl KafkaConnection {

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let r = self.stream.write(&msg[..]).map_err(From::from);
        debug!("Sent {} bytes to: {:?} => {:?}", msg.len(), self, r);
        r
    }

    pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
        let r = (&mut self.stream).read_exact(buf).map_err(From::from);
        debug!("Read {} bytes from: {:?} => {:?}", buf.len(), self, r);
        r
    }

    pub fn read_exact_alloc(&mut self, size: u64) -> Result<Vec<u8>> {
        let size: usize = size as usize;
        let mut buffer: Vec<u8> = Vec::with_capacity(size);
        // this is safe actually: we are setting the len to the
        // buffers capacity and either fully populate it in the
        // following call to `read_exact` or discard the vector (in
        // the error case)
        unsafe { buffer.set_len(size) };
        try!(self.read_exact(buffer.as_mut_slice()));
        Ok(buffer)
    }

    fn shutdown(&mut self) -> Result<()> {
        let r = self.stream.shutdown(Shutdown::Both);
        debug!("Shut down: {:?} => {:?}", self, r);
        r.map_err(From::from)
    }

    fn from_stream(stream: KafkaStream, id: u32, host: &str, rw_timeout: Option<Duration>)
                   -> Result<KafkaConnection>
    {
        try!(stream.set_read_timeout(rw_timeout));
        try!(stream.set_write_timeout(rw_timeout));
        Ok(KafkaConnection { id: id, host: host.to_owned(), stream: stream })
    }

    #[cfg(not(feature = "security"))]
    fn new(id: u32, host: &str, rw_timeout: Option<Duration>) -> Result<KafkaConnection> {
        KafkaConnection::from_stream(try!(TcpStream::connect(host)), id, host, rw_timeout)
    }

    #[cfg(feature = "security")]
    fn new(id: u32, host: &str, rw_timeout: Option<Duration>, security: Option<&SslContext>)
           -> Result<KafkaConnection>
    {
        let stream = try!(TcpStream::connect(host));
        let stream = match security {
            Some(ctx) => KafkaStream::Ssl(try!(SslStream::connect(ctx, stream))),
            None => KafkaStream::Plain(stream),
        };
        KafkaConnection::from_stream(stream, id, host, rw_timeout)
    }
}
<commit_msg>Lower log level for read/send network debug messages<commit_after>//! Network related functionality for `KafkaClient`.
//!
//! This module is crate private and not exposed to the public except
//! through re-exports of individual items from within
//! `kafka::client`.

use std::collections::HashMap;
use std::fmt;
use std::io::{Read, Write};
use std::mem;
use std::net::{TcpStream, Shutdown};
use std::time::{Instant, Duration};

#[cfg(feature = "security")]
use openssl::ssl::{SslContext, SslStream};

use error::Result;

// --------------------------------------------------------------------

/// Security relevant configuration options for `KafkaClient`.
// This will be expanded in the future. See #51.
#[cfg(feature = "security")]
#[derive(Debug)]
pub struct SecurityConfig(SslContext);

#[cfg(feature = "security")]
impl SecurityConfig {
    /// In the future this will also support a kerbos via #51.
    pub fn new(ssl: SslContext) -> SecurityConfig {
        SecurityConfig(ssl)
    }
}

// --------------------------------------------------------------------

struct Pooled<T> {
    last_checkout: Instant,
    item: T,
}

impl<T> Pooled<T> {
    fn new(last_checkout: Instant, item: T) -> Self {
        Pooled { last_checkout: last_checkout, item: item }
    }
}

impl<T: fmt::Debug> fmt::Debug for Pooled<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Pooled {{ last_checkout: {:?}, item: {:?} }}",
               self.last_checkout, self.item)
    }
}

#[derive(Debug)]
pub struct Config {
    rw_timeout: Option<Duration>,
    idle_timeout: Duration,
    #[cfg(feature = "security")]
    security_config: Option<SecurityConfig>,
}

impl Config {
    #[cfg(not(feature = "security"))]
    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
        KafkaConnection::new(id, host, self.rw_timeout)
            .map(|c| {
                debug!("Established: {:?}", c);
                c
            })
    }

    #[cfg(feature = "security")]
    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
        KafkaConnection::new(id,
                             host,
                             self.rw_timeout,
                             self.security_config.as_ref().map(|c| &c.0))
            .map(|c| {
                debug!("Established: {:?}", c);
                c
            })
    }
}

#[derive(Debug)]
struct State {
    num_conns: u32,
}

impl State {
    fn new() -> State {
        State { num_conns: 0 }
    }

    fn next_conn_id(&mut self) -> u32 {
        let c = self.num_conns;
        self.num_conns = self.num_conns.wrapping_add(1);
        c
    }
}

#[derive(Debug)]
pub struct Connections {
    conns: HashMap<String, Pooled<KafkaConnection>>,
    state: State,
    config: Config,
}

impl Connections {
    #[cfg(not(feature = "security"))]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
        Connections {
            conns: HashMap::new(),
            state: State::new(),
            config: Config {
                rw_timeout: rw_timeout,
                idle_timeout: idle_timeout,
            },
        }
    }

    #[cfg(feature = "security")]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
        Self::new_with_security(rw_timeout, idle_timeout, None)
    }

    #[cfg(feature = "security")]
    pub fn new_with_security(rw_timeout: Option<Duration>,
                             idle_timeout: Duration,
                             security: Option<SecurityConfig>)
                             -> Connections
    {
        Connections {
            conns: HashMap::new(),
            state: State::new(),
            config: Config {
                rw_timeout: rw_timeout,
                idle_timeout: idle_timeout,
                security_config: security,
            },
        }
    }

    pub fn set_idle_timeout(&mut self, idle_timeout: Duration) {
        self.config.idle_timeout = idle_timeout;
    }

    pub fn idle_timeout(&self) -> Duration {
        self.config.idle_timeout
    }

    pub fn get_conn<'a>(&'a mut self, host: &str, now: Instant)
                        -> Result<&'a mut KafkaConnection>
    {
        if let Some(conn) = self.conns.get_mut(host) {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout reached: {:?}", conn.item);
                let new_conn = try!(self.config.new_conn(self.state.next_conn_id(), host));
                let _ = conn.item.shutdown();
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            // ~ decouple the lifetimes to make the borrowck happy;
            // this is safe since we're immediatelly returning the
            // reference and the rest of the code in this method is
            // not affected
            return Ok(unsafe { mem::transmute(kconn) });
        }
        let cid = self.state.next_conn_id();
        self.conns.insert(host.to_owned(),
                          Pooled::new(now, try!(self.config.new_conn(cid, host))));
        Ok(&mut self.conns.get_mut(host).unwrap().item)
    }

    pub fn get_conn_any(&mut self, now: Instant) -> Option<&mut KafkaConnection> {
        for (host, conn) in &mut self.conns {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout reached: {:?}", conn.item);
                let new_conn_id = self.state.next_conn_id();
                let new_conn = match self.config.new_conn(new_conn_id, host.as_str()) {
                    Ok(new_conn) => {
                        let _ = conn.item.shutdown();
                        new_conn
                    }
                    Err(e) => {
                        warn!("Failed to establish connection to {}: {:?}", host, e);
                        continue;
                    }
                };
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            return Some(kconn);
        }
        None
    }
}

// --------------------------------------------------------------------

trait IsSecured {
    fn is_secured(&self) -> bool;
}

#[cfg(not(feature = "security"))]
type KafkaStream = TcpStream;

#[cfg(not(feature = "security"))]
impl IsSecured for KafkaStream {
    fn is_secured(&self) -> bool {
        false
    }
}

#[cfg(feature = "security")]
use self::openssled::KafkaStream;

#[cfg(feature = "security")]
mod openssled {
    use std::io::{self, Read, Write};
    use std::net::{TcpStream, Shutdown};
    use std::time::Duration;

    use openssl::ssl::SslStream;

    use super::IsSecured;

    pub enum KafkaStream {
        Plain(TcpStream),
        Ssl(SslStream<TcpStream>)
    }

    impl IsSecured for KafkaStream {
        fn is_secured(&self) -> bool {
            match *self {
                KafkaStream::Ssl(_) => true,
                _ => false,
            }
        }
    }

    impl KafkaStream {
        fn get_ref(&self) -> &TcpStream {
            match *self {
                KafkaStream::Plain(ref s) => s,
                KafkaStream::Ssl(ref s) => s.get_ref()
            }
        }

        pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_read_timeout(dur)
        }

        pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_write_timeout(dur)
        }

        pub fn shutdown(&mut self, how: Shutdown) -> io::Result<()> {
            self.get_ref().shutdown(how)
        }
    }

    impl Read for KafkaStream {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.read(buf),
                KafkaStream::Ssl(ref mut s) => s.read(buf),
            }
        }
    }

    impl Write for KafkaStream {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.write(buf),
                KafkaStream::Ssl(ref mut s) => s.write(buf),
            }
        }
        fn flush(&mut self) -> io::Result<()> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.flush(),
                KafkaStream::Ssl(ref mut s) => s.flush(),
            }
        }
    }
}

/// A TCP stream to a remote Kafka broker.
pub struct KafkaConnection {
    // a surrogate identifier to distinguish between
    // connections to the same host in debug messages
    id: u32,
    // "host:port"
    host: String,
    // the (wrapped) tcp stream
    stream: KafkaStream,
}

impl fmt::Debug for KafkaConnection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "KafkaConnection {{ id: {}, secured: {}, host: \"{}\" }}",
               self.id, self.stream.is_secured(), self.host)
    }
}

impl KafkaConnection {

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let r = self.stream.write(&msg[..]).map_err(From::from);
        trace!("Sent {} bytes to: {:?} => {:?}", msg.len(), self, r);
        r
    }

    pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
        let r = (&mut self.stream).read_exact(buf).map_err(From::from);
        trace!("Read {} bytes from: {:?} => {:?}", buf.len(), self, r);
        r
    }

    pub fn read_exact_alloc(&mut self, size: u64) -> Result<Vec<u8>> {
        let size: usize = size as usize;
        let mut buffer: Vec<u8> = Vec::with_capacity(size);
        // this is safe actually: we are setting the len to the
        // buffers capacity and either fully populate it in the
        // following call to `read_exact` or discard the vector (in
        // the error case)
        unsafe { buffer.set_len(size) };
        try!(self.read_exact(buffer.as_mut_slice()));
        Ok(buffer)
    }

    fn shutdown(&mut self) -> Result<()> {
        let r = self.stream.shutdown(Shutdown::Both);
        debug!("Shut down: {:?} => {:?}", self, r);
        r.map_err(From::from)
    }

    fn from_stream(stream: KafkaStream, id: u32, host: &str, rw_timeout: Option<Duration>)
                   -> Result<KafkaConnection>
    {
        try!(stream.set_read_timeout(rw_timeout));
        try!(stream.set_write_timeout(rw_timeout));
        Ok(KafkaConnection { id: id, host: host.to_owned(), stream: stream })
    }

    #[cfg(not(feature = "security"))]
    fn new(id: u32, host: &str, rw_timeout: Option<Duration>) -> Result<KafkaConnection> {
        KafkaConnection::from_stream(try!(TcpStream::connect(host)), id, host, rw_timeout)
    }

    #[cfg(feature = "security")]
    fn new(id: u32, host: &str, rw_timeout: Option<Duration>, security: Option<&SslContext>)
           -> Result<KafkaConnection>
    {
        let stream = try!(TcpStream::connect(host));
        let stream = match security {
            Some(ctx) => KafkaStream::Ssl(try!(SslStream::connect(ctx, stream))),
            None => KafkaStream::Plain(stream),
        };
        KafkaConnection::from_stream(stream, id, host, rw_timeout)
    }
}
<|endoftext|>


<commit_before><commit_msg>Added a call method to operations::expression to make it easier for prelude functions to call other methods<commit_after><|endoftext|>


<commit_before>//! The `arrange` operator and its variants arrange a collection into a shareable trace structure.
//!
//! The `arrange` operator is applied to a differential dataflow `Collection` and return an `Arranged`.
//! Several operators (`join`, `group`, and `cogroup`, among others) are implemented in terms of this
//! structure, and can be applied directly to arranged data. Internally, the operators will borrow the
//! shared state, and listen on the timely stream for shared batches of data. The resources to assemble
//! the collection, both computation and memory, are spent only once.
//!
//! The arranged collection is stored in a trace, whose append-only behavior means that it is safe to
//! share between the single writer and multiple readers. Each reader is expected to interrogate the
//! trace only at times for which it knows the trace is complete, as indicated by the frontiers on its
//! incoming channels. Failing to do this is "safe", but may result in undefined semantics.
//!
//! Internally, the shared trace is wrapped in a `TraceWrapper` type which maintains information about
//! the frontiers of all of its referees. Each referee has a `TraceHandle`, which acts as a reference
//! counted pointer, and which mediates the advancement of frontiers. Ideally, a `TraceHandle` looks a
//! lot like a trace, though this isn't beatifully masked at the moment (it can't implement the trait
//! because we can't insert at it; it does implement `advance_by` and could implement `cursor`).

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use linear_map::LinearMap;

use timely::dataflow::*;
use timely::dataflow::operators::Unary;
use timely::dataflow::channels::pact::Exchange;
use timely::progress::frontier::MutableAntichain;

use ::{Data, Collection};
use lattice::Lattice;
use trace::{Trace, Batch, Builder};

use trace::implementations::trie::Spine as TrieSpine;
use trace::implementations::keys::Spine as KeysSpine;

/// Wrapper type to permit transfer of `Rc` types, as in batch.
#[derive(Clone)]
pub struct BatchWrapper<T> {
    /// The wrapped item.
    pub item: T,
}

// NOTE: This is all horrible. Don't look too hard.
impl<T> ::abomonation::Abomonation for BatchWrapper<T> {
   unsafe fn entomb(&self, _writer: &mut Vec<u8>) { panic!() }
   unsafe fn embalm(&mut self) { panic!() }
   unsafe fn exhume<'a,'b>(&'a mut self, _bytes: &'b mut [u8]) -> Option<&'b mut [u8]> { panic!()  }
}

/// A wrapper around a trace which tracks the frontiers of all referees.
pub struct TraceWrapper<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> {
    phantom: ::std::marker::PhantomData<(K, V)>,
    frontiers: MutableAntichain<T>,
    /// The wrapped trace.
    pub trace: Tr,
}

impl<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> TraceWrapper<K, V, T, Tr> {
    /// Allocates a new trace wrapper.
    pub fn new(default: T) -> Self {
        TraceWrapper {
            phantom: ::std::marker::PhantomData,
            frontiers: MutableAntichain::new(),
            trace: Tr::new(default)
        }
    }
    /// Reports the current frontier of the trace.
    fn frontier(&self) -> &[T] { self.frontiers.elements() }
    /// Replaces elements of `lower` with those of `upper`.
    fn adjust_frontier(&mut self, lower: &[T], upper: &[T]) {
        for element in upper { self.frontiers.update_and(element, 1, |_,_| {}); }
        for element in lower { self.frontiers.update_and(element, -1, |_,_| {}); }
        self.trace.advance_by(self.frontiers.elements());
    }
}

/// A handle to a shared trace which maintains its own frontier information.
///
/// As long as the handle exists, it should protect the trace from advancing past the associated frontier.
/// This protection advances as the handle's frontier is advanced. When the handle is dropped the protection
/// is removed.
pub struct TraceHandle<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> {
    phantom: ::std::marker::PhantomData<(K, V)>,
    frontier: Vec<T>,
    /// Wrapped trace. Please be gentle when using.
    pub wrapper: Rc<RefCell<TraceWrapper<K, V, T, Tr>>>,
}

impl<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> TraceHandle<K, V, T, Tr> {
    /// Allocates a new handle from an existing wrapped wrapper.
    pub fn new(wrapper: &Rc<RefCell<TraceWrapper<K, V, T, Tr>>>) -> Self {
        let frontier = wrapper.borrow().frontier().to_vec();
        wrapper.borrow_mut().adjust_frontier(&[], &frontier[..]);
        TraceHandle {
            phantom: ::std::marker::PhantomData,
            frontier: frontier,
            wrapper: wrapper.clone(),
        }
    }
    /// Sets frontier to now be elements in `frontier`.
    pub fn advance_by(&mut self, frontier: &[T]) {
        self.wrapper.borrow_mut().adjust_frontier(&self.frontier[..], frontier);
        self.frontier = frontier.to_vec();
    }
    /// Creates a new cursor over the wrapped trace.
    pub fn cursor(&self) -> Tr::Cursor {
        self.wrapper.borrow().trace.cursor()
    }
}

impl<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> Drop for TraceHandle<K, V, T, Tr> {
    fn drop(&mut self) {
        self.wrapper.borrow_mut().adjust_frontier(&self.frontier[..], &[]);
        self.frontier = Vec::new();
    }
}

/// A collection of `(K,V)` values as a timely stream and shared trace.
///
/// An `Arranged` performs the task of arranging a keyed collection once,
/// allowing multiple differential operators to use the same trace. This
/// saves on computation and memory, in exchange for some cognitive overhead
/// in writing differential operators: each must pay enough care to signals
/// from the `stream` field to know the subset of `trace` it has logically
/// received.
pub struct Arranged<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>>
    where G::Timestamp: Lattice+Ord {
    /// A stream containing arranged updates.
    ///
    /// This stream may be reduced to just the keys, with the expectation that looking up the
    /// updated values in the associated trace should be more efficient than cloning the values.
    /// Users who need the values can obviously use the original stream.
    pub stream: Stream<G, BatchWrapper<T::Batch>>,
    /// A shared trace, updated by the `Arrange` operator and readable by others.
    pub trace: Rc<RefCell<TraceWrapper<K, V, G::Timestamp, T>>>,
    /// If dereferenced, we build this collection.
    pub collection: Option<Collection<G, (K, V)>>,
}

impl<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>> Arranged<G, K, V, T>
    where G::Timestamp: Lattice+Ord,
          T::Batch: Clone+'static {

    /// Allocates a new handle to the shared trace, with independent frontier tracking.
    pub fn new_handle(&self) -> TraceHandle<K, V, G::Timestamp, T> {
        TraceHandle::new(&self.trace)
    }

    /// Flattens the stream into a `Collection`.
    ///
    /// This operator is not obviously more efficient than using the `Collection` that is input
    /// to the `Arrange` operator. However, this may be the only way to recover a collection from
    /// only the `Arranged` type.
    pub fn as_collection(&self) -> Collection<G, (K, V)> {
        unimplemented!()
    }
}

// impl<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>> Deref for Arranged<G, K, V, T>
//     where G::Timestamp: Lattice+Ord,
//           T::Batch: Clone+'static {
//     type Target = Collection<G, (K, V)>;
//     fn deref(&self) -> &Self::Target {
//         if self.collection.is_none() {
//             self.collection = Some(self.as_collection());
//         }
//         self.collection.as_ref().unwrap()
//     }
// }

/// Arranges something as `(Key,Val)` pairs according to a type `T` of trace.
pub trait Arrange<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>>
where G::Timestamp: Lattice+Ord,
      T::Batch: Clone+'static {
    /// Arranges a stream of `(Key, Val)` updates by `Key`.
    ///
    /// This operator arranges a stream of values into a shared trace, whose contents it maintains.
    /// This trace is current for all times completed by the output stream, which can be used to
    /// safely identify the stable times and values in the trace.
    fn arrange(&self) -> Arranged<G, K, V, T>;
}

impl<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>+'static> Arrange<G, K, V, T> for Collection<G, (K, V)>
where G::Timestamp: Lattice+Ord,
      T::Batch: Clone+'static {
    fn arrange(&self) -> Arranged<G, K, V, T> {

        // create a trace to share with downstream consumers.
        let trace = Rc::new(RefCell::new(TraceWrapper::new(Default::default())));
        let source = Rc::downgrade(&trace);

        // A map from times to received (key, val, wgt) triples.
        let mut inputs = LinearMap::new();

        // fabricate a data-parallel operator using the `unary_notify` pattern.
        let exchange = Exchange::new(|x: &((K,V),isize)| (x.0).0.hashed());
        let stream = self.inner.unary_notify(exchange, "ArrangeByKey", vec![], move |input, output, notificator| {

            input.for_each(|time, data| {
                inputs.entry(time.time())
                      .or_insert_with(|| { notificator.notify_at(time.clone()); <T::Batch as Batch<K,V,G::Timestamp>>::Builder::new() })
                      .extend(data.drain(..).map(|(kv,d)| (kv.0, kv.1, time.time(), d)));
            });

            notificator.for_each(|index, _count, _notificator| {
                if let Some(builder) = inputs.remove(&index) {
                    let batch = builder.done(&[], &[]);
                    source.upgrade().map(|trace| {
                        let trace: &mut T = &mut trace.borrow_mut().trace;
                        trace.insert(batch.clone())
                    });
                    output.session(&index).give(BatchWrapper { item: batch });
                }
            });
        });

        Arranged { stream: stream, trace: trace, collection: None }
    }
}

/// Arranges something as `(Key,Val)` pairs according to a type `T` of trace.
pub trait ArrangeByKey<G: Scope, K: Data+Ord, V: Data+Ord>
where G::Timestamp: Lattice+Ord {
    /// Arranges a stream of `(Key, Val)` updates by `Key`.
    ///
    /// This operator arranges a stream of values into a shared trace, whose contents it maintains.
    /// This trace is current for all times completed by the output stream, which can be used to
    /// safely identify the stable times and values in the trace.
    fn arrange_by_key(&self) -> Arranged<G, K, V, TrieSpine<K, V, G::Timestamp>>;
}

impl<G: Scope, K: Data+Ord, V: Data+Ord> ArrangeByKey<G, K, V> for Collection<G, (K,V)>
where G::Timestamp: Lattice+Ord  {

    fn arrange_by_key(&self) -> Arranged<G, K, V, TrieSpine<K, V, G::Timestamp>> {
        self.arrange()
    }
}

/// Arranges something as `(Key,Val)` pairs according to a type `T` of trace.
pub trait ArrangeBySelf<G: Scope, K: Data+Ord>
where G::Timestamp: Lattice+Ord {
    /// Arranges a stream of `(Key, Val)` updates by `Key`.
    ///
    /// This operator arranges a stream of values into a shared trace, whose contents it maintains.
    /// This trace is current for all times completed by the output stream, which can be used to
    /// safely identify the stable times and values in the trace.
    fn arrange_by_self(&self) -> Arranged<G, K, (), KeysSpine<K, G::Timestamp>>;
}

impl<G: Scope, K: Data+Ord> ArrangeBySelf<G, K> for Collection<G, K>
where G::Timestamp: Lattice+Ord {
    fn arrange_by_self(&self) -> Arranged<G, K, (), KeysSpine<K, G::Timestamp>> {
        self.map(|k| (k,()))
            .arrange()
    }
}
<commit_msg>added ascollection<commit_after>//! The `arrange` operator and its variants arrange a collection into a shareable trace structure.
//!
//! The `arrange` operator is applied to a differential dataflow `Collection` and return an `Arranged`.
//! Several operators (`join`, `group`, and `cogroup`, among others) are implemented in terms of this
//! structure, and can be applied directly to arranged data. Internally, the operators will borrow the
//! shared state, and listen on the timely stream for shared batches of data. The resources to assemble
//! the collection, both computation and memory, are spent only once.
//!
//! The arranged collection is stored in a trace, whose append-only behavior means that it is safe to
//! share between the single writer and multiple readers. Each reader is expected to interrogate the
//! trace only at times for which it knows the trace is complete, as indicated by the frontiers on its
//! incoming channels. Failing to do this is "safe", but may result in undefined semantics.
//!
//! Internally, the shared trace is wrapped in a `TraceWrapper` type which maintains information about
//! the frontiers of all of its referees. Each referee has a `TraceHandle`, which acts as a reference
//! counted pointer, and which mediates the advancement of frontiers. Ideally, a `TraceHandle` looks a
//! lot like a trace, though this isn't beatifully masked at the moment (it can't implement the trait
//! because we can't insert at it; it does implement `advance_by` and could implement `cursor`).
//!
//! Note: in the current implementation, all batches have only one time. This will probably change and
//! shouldn't be relied on.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use linear_map::LinearMap;

use timely::dataflow::*;
use timely::dataflow::operators::Unary;
use timely::dataflow::channels::pact::{Pipeline, Exchange};
use timely::progress::frontier::MutableAntichain;

use ::{Data, Collection, AsCollection};
use lattice::Lattice;
use trace::{Trace, Batch, Builder, Cursor};
use trace::implementations::trie::Spine as TrieSpine;
use trace::implementations::keys::Spine as KeysSpine;

/// Wrapper type to permit transfer of `Rc` types, as in batch.
#[derive(Clone)]
pub struct BatchWrapper<T> {
    /// The wrapped item.
    pub item: T,
}

// NOTE: This is all horrible. Don't look too hard.
impl<T> ::abomonation::Abomonation for BatchWrapper<T> {
   unsafe fn entomb(&self, _writer: &mut Vec<u8>) { panic!() }
   unsafe fn embalm(&mut self) { panic!() }
   unsafe fn exhume<'a,'b>(&'a mut self, _bytes: &'b mut [u8]) -> Option<&'b mut [u8]> { panic!()  }
}

/// A wrapper around a trace which tracks the frontiers of all referees.
pub struct TraceWrapper<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> {
    phantom: ::std::marker::PhantomData<(K, V)>,
    frontiers: MutableAntichain<T>,
    /// The wrapped trace.
    pub trace: Tr,
}

impl<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> TraceWrapper<K, V, T, Tr> {
    /// Allocates a new trace wrapper.
    pub fn new(default: T) -> Self {
        TraceWrapper {
            phantom: ::std::marker::PhantomData,
            frontiers: MutableAntichain::new(),
            trace: Tr::new(default)
        }
    }
    /// Reports the current frontier of the trace.
    fn frontier(&self) -> &[T] { self.frontiers.elements() }
    /// Replaces elements of `lower` with those of `upper`.
    fn adjust_frontier(&mut self, lower: &[T], upper: &[T]) {
        for element in upper { self.frontiers.update_and(element, 1, |_,_| {}); }
        for element in lower { self.frontiers.update_and(element, -1, |_,_| {}); }
        self.trace.advance_by(self.frontiers.elements());
    }
}

/// A handle to a shared trace which maintains its own frontier information.
///
/// As long as the handle exists, it should protect the trace from advancing past the associated frontier.
/// This protection advances as the handle's frontier is advanced. When the handle is dropped the protection
/// is removed.
pub struct TraceHandle<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> {
    phantom: ::std::marker::PhantomData<(K, V)>,
    frontier: Vec<T>,
    /// Wrapped trace. Please be gentle when using.
    pub wrapper: Rc<RefCell<TraceWrapper<K, V, T, Tr>>>,
}

impl<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> TraceHandle<K, V, T, Tr> {
    /// Allocates a new handle from an existing wrapped wrapper.
    pub fn new(wrapper: &Rc<RefCell<TraceWrapper<K, V, T, Tr>>>) -> Self {
        let frontier = wrapper.borrow().frontier().to_vec();
        wrapper.borrow_mut().adjust_frontier(&[], &frontier[..]);
        TraceHandle {
            phantom: ::std::marker::PhantomData,
            frontier: frontier,
            wrapper: wrapper.clone(),
        }
    }
    /// Sets frontier to now be elements in `frontier`.
    pub fn advance_by(&mut self, frontier: &[T]) {
        self.wrapper.borrow_mut().adjust_frontier(&self.frontier[..], frontier);
        self.frontier = frontier.to_vec();
    }
    /// Creates a new cursor over the wrapped trace.
    pub fn cursor(&self) -> Tr::Cursor {
        self.wrapper.borrow().trace.cursor()
    }
}

impl<K: Ord, V: Ord, T: Lattice+Ord+Clone+::std::fmt::Debug+'static, Tr: Trace<K, V, T>> Drop for TraceHandle<K, V, T, Tr> {
    fn drop(&mut self) {
        self.wrapper.borrow_mut().adjust_frontier(&self.frontier[..], &[]);
        self.frontier = Vec::new();
    }
}

/// A collection of `(K,V)` values as a timely stream and shared trace.
///
/// An `Arranged` performs the task of arranging a keyed collection once,
/// allowing multiple differential operators to use the same trace. This
/// saves on computation and memory, in exchange for some cognitive overhead
/// in writing differential operators: each must pay enough care to signals
/// from the `stream` field to know the subset of `trace` it has logically
/// received.
pub struct Arranged<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>>
    where G::Timestamp: Lattice+Ord {
    /// A stream containing arranged updates.
    ///
    /// This stream may be reduced to just the keys, with the expectation that looking up the
    /// updated values in the associated trace should be more efficient than cloning the values.
    /// Users who need the values can obviously use the original stream.
    pub stream: Stream<G, BatchWrapper<T::Batch>>,
    /// A shared trace, updated by the `Arrange` operator and readable by others.
    pub trace: Rc<RefCell<TraceWrapper<K, V, G::Timestamp, T>>>,
    /// If dereferenced, we build this collection.
    pub collection: Option<Collection<G, (K, V)>>,
}

impl<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>> Arranged<G, K, V, T>
    where G::Timestamp: Lattice+Ord,
          T::Batch: Clone+'static {

    /// Allocates a new handle to the shared trace, with independent frontier tracking.
    pub fn new_handle(&self) -> TraceHandle<K, V, G::Timestamp, T> {
        TraceHandle::new(&self.trace)
    }

    /// Flattens the stream into a `Collection`.
    ///
    /// This operator is not obviously more efficient than using the `Collection` that is input
    /// to the `Arrange` operator. However, this may be the only way to recover a collection from
    /// only the `Arranged` type.
    pub fn as_collection(&self) -> Collection<G, (K, V)> {

        // unimplemented!()

        self.stream.unary_stream(Pipeline, "AsCollection", |input, output| {
            input.for_each(|time, data| {
                let mut session = output.session(&time);
                for wrapper in data.drain(..) {
                    let batch = wrapper.item;
                    let mut cursor = batch.cursor();
                    while cursor.key_valid() {
                        while cursor.val_valid() {
                            cursor.map_times(|_time, diff| {
                                session.give(((cursor.key().clone(), cursor.val().clone()), diff));
                            });
                            cursor.step_val();
                        }
                        cursor.step_key();
                    }
                }
            })
        })
        .as_collection()
    }
}

// impl<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>> Deref for Arranged<G, K, V, T>
//     where G::Timestamp: Lattice+Ord,
//           T::Batch: Clone+'static {
//     type Target = Collection<G, (K, V)>;
//     fn deref(&self) -> &Self::Target {
//         if self.collection.is_none() {
//             self.collection = Some(self.as_collection());
//         }
//         self.collection.as_ref().unwrap()
//     }
// }

/// Arranges something as `(Key,Val)` pairs according to a type `T` of trace.
pub trait Arrange<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>>
where G::Timestamp: Lattice+Ord,
      T::Batch: Clone+'static {
    /// Arranges a stream of `(Key, Val)` updates by `Key`.
    ///
    /// This operator arranges a stream of values into a shared trace, whose contents it maintains.
    /// This trace is current for all times completed by the output stream, which can be used to
    /// safely identify the stable times and values in the trace.
    fn arrange(&self) -> Arranged<G, K, V, T>;
}

impl<G: Scope, K: Data+Ord, V: Data+Ord, T: Trace<K, V, G::Timestamp>+'static> Arrange<G, K, V, T> for Collection<G, (K, V)>
where G::Timestamp: Lattice+Ord,
      T::Batch: Clone+'static {
    fn arrange(&self) -> Arranged<G, K, V, T> {

        // create a trace to share with downstream consumers.
        let trace = Rc::new(RefCell::new(TraceWrapper::new(Default::default())));
        let source = Rc::downgrade(&trace);

        // A map from times to received (key, val, wgt) triples.
        let mut inputs = LinearMap::new();

        // fabricate a data-parallel operator using the `unary_notify` pattern.
        let exchange = Exchange::new(|x: &((K,V),isize)| (x.0).0.hashed());
        let stream = self.inner.unary_notify(exchange, "ArrangeByKey", vec![], move |input, output, notificator| {

            input.for_each(|time, data| {
                inputs.entry(time.time())
                      .or_insert_with(|| { notificator.notify_at(time.clone()); <T::Batch as Batch<K,V,G::Timestamp>>::Builder::new() })
                      .extend(data.drain(..).map(|(kv,d)| (kv.0, kv.1, time.time(), d)));
            });

            notificator.for_each(|index, _count, _notificator| {
                if let Some(builder) = inputs.remove(&index) {
                    let batch = builder.done(&[], &[]);
                    source.upgrade().map(|trace| {
                        let trace: &mut T = &mut trace.borrow_mut().trace;
                        trace.insert(batch.clone())
                    });
                    output.session(&index).give(BatchWrapper { item: batch });
                }
            });
        });

        Arranged { stream: stream, trace: trace, collection: None }
    }
}

/// Arranges something as `(Key,Val)` pairs according to a type `T` of trace.
pub trait ArrangeByKey<G: Scope, K: Data+Ord, V: Data+Ord>
where G::Timestamp: Lattice+Ord {
    /// Arranges a stream of `(Key, Val)` updates by `Key`.
    ///
    /// This operator arranges a stream of values into a shared trace, whose contents it maintains.
    /// This trace is current for all times completed by the output stream, which can be used to
    /// safely identify the stable times and values in the trace.
    fn arrange_by_key(&self) -> Arranged<G, K, V, TrieSpine<K, V, G::Timestamp>>;
}

impl<G: Scope, K: Data+Ord, V: Data+Ord> ArrangeByKey<G, K, V> for Collection<G, (K,V)>
where G::Timestamp: Lattice+Ord  {

    fn arrange_by_key(&self) -> Arranged<G, K, V, TrieSpine<K, V, G::Timestamp>> {
        self.arrange()
    }
}

/// Arranges something as `(Key,Val)` pairs according to a type `T` of trace.
pub trait ArrangeBySelf<G: Scope, K: Data+Ord>
where G::Timestamp: Lattice+Ord {
    /// Arranges a stream of `(Key, Val)` updates by `Key`.
    ///
    /// This operator arranges a stream of values into a shared trace, whose contents it maintains.
    /// This trace is current for all times completed by the output stream, which can be used to
    /// safely identify the stable times and values in the trace.
    fn arrange_by_self(&self) -> Arranged<G, K, (), KeysSpine<K, G::Timestamp>>;
}

impl<G: Scope, K: Data+Ord> ArrangeBySelf<G, K> for Collection<G, K>
where G::Timestamp: Lattice+Ord {
    fn arrange_by_self(&self) -> Arranged<G, K, (), KeysSpine<K, G::Timestamp>> {
        self.map(|k| (k,()))
            .arrange()
    }
}
<|endoftext|>


<commit_before>// Copyright 2019-2020 Google LLC
//
// 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.

//! Helps manipulate bit fields in 32-bits words.
// TODO(ia0): Remove when the module is used.
#![cfg_attr(not(test), allow(dead_code, unused_macros))]

use crate::{StoreError, StoreResult};

/// Represents a bit field.
///
/// A bit field is a contiguous sequence of bits in a 32-bits word.
///
/// # Invariant
///
/// - `pos + len < 32`.
pub struct Field {
    /// The position of the bit field.
    pub pos: usize,

    /// The length of the bit field.
    pub len: usize,
}

impl Field {
    /// Reads the value of a bit field.
    pub fn get(&self, word: u32) -> usize {
        ((word >> self.pos) & self.mask()) as usize
    }

    /// Sets the value of a bit field.
    ///
    /// # Preconditions
    ///
    /// - `num_bits(value) < self.len`.
    /// - `self.get(*word) & value == value`.
    pub fn set(&self, word: &mut u32, value: usize) {
        let value = value as u32;
        debug_assert_eq!(value & self.mask(), value);
        let mask = !(self.mask() << self.pos);
        *word &= mask | (value << self.pos);
        debug_assert_eq!(self.get(*word), value as usize);
    }

    /// Returns a bit mask the length of the bit field.
    fn mask(&self) -> u32 {
        (1 << self.len) - 1
    }
}

/// Represents a constant bit field.
///
/// # Invariant
///
/// - `num_bits(value) <= field.len`.
pub struct Const {
    /// The bit field.
    pub field: Field,

    /// The constant value.
    pub value: usize,
}

impl Const {
    /// Checks that the bit field has its value.
    pub fn check(&self, word: u32) -> bool {
        self.field.get(word) == self.value
    }

    /// Sets the bit field to its value.
    pub fn set(&self, word: &mut u32) {
        self.field.set(word, self.value);
    }
}

/// Represents a single bit.
///
/// # Invariant
///
/// - `pos < 32`.
pub struct Bit {
    /// The position of the bit.
    pub pos: usize,
}

impl Bit {
    /// Returns whether the value of the bit is zero.
    pub fn get(&self, word: u32) -> bool {
        word & (1 << self.pos) == 0
    }

    /// Sets the value of the bit to zero.
    pub fn set(&self, word: &mut u32) {
        *word &= !(1 << self.pos);
    }
}

/// Represents a checksum.
///
/// A checksum is a bit field counting how many bits are set to zero in the word (except in the
/// checksum itself) plus some external increment. It essentially behaves like a bit field storing
/// the external increment.
pub struct Checksum {
    /// The bit field
    pub field: Field,
}

impl Checksum {
    /// Reads the external increment from the checksum.
    ///
    /// # Errors
    ///
    /// Returns `InvalidStorage` if the external increment would be negative.
    pub fn get(&self, word: u32) -> StoreResult<usize> {
        let checksum = self.field.get(word);
        let zeros = word.count_zeros() as usize - (self.field.len - checksum.count_ones() as usize);
        checksum
            .checked_sub(zeros)
            .ok_or(StoreError::InvalidStorage)
    }

    /// Sets the checksum to the external increment value.
    ///
    /// # Preconditions
    ///
    /// - The checksum bits should be set to one.
    /// - The number of zeros in the word plus the value should fit in the checksum.
    pub fn set(&self, word: &mut u32, value: usize) {
        debug_assert_eq!(self.field.get(*word), self.field.mask() as usize);
        self.field.set(word, word.count_zeros() as usize + value);
    }
}

/// Tracks the number of bits used so far.
///
/// # Features
///
/// Only available for tests.
#[cfg(any(doc, test))]
pub struct Length {
    /// The position of the next available bit.
    pub pos: usize,
}

/// Helps defining contiguous bit fields.
///
/// It takes a sequence of bit field descriptors as argument. A bit field descriptor is one of the
/// following:
/// - `$name: Bit,` to define a bit
/// - `$name: Field <= $max,` to define a bit field of minimum length to store `$max`
/// - `$name: Checksum <= $max,` to define a checksum of minimum length to store `$max`
/// - `$name: Length,` to define a length tracker
/// - `$name: Const = [$bits],` to define a constant bit field with value `$bits` (a sequence of
///   space-separated bits)
#[cfg_attr(doc, macro_export)] // For `cargo doc` to produce documentation.
macro_rules! bitfield {
    ($($input: tt)*) => {
        bitfield_impl! { []{ pos: 0 }[$($input)*] }
    };
}

macro_rules! bitfield_impl {
    // Main rules:
    // - Input are bit field descriptors
    // - Position is the number of bits used by prior bit fields
    // - Output are the bit field definitions
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Bit, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Bit = Bit { pos: $pos };]
            { pos: $pos + 1 }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Field <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Field = Field { pos: $pos, len: num_bits($max) };]
            { pos: $pos + $name.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$name: ident: Checksum <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Checksum = Checksum {
                field: Field { pos: $pos, len: num_bits($max) }
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$(#[$meta: meta])* $name: ident: Length, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* $(#[$meta])* const $name: Length = Length { pos: $pos };]
            { pos: $pos }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Const = $bits: tt, $($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name []$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[]) => { $($output)* };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits
    // - Output is the reversed sequence of bits
    (Reverse $name: ident [$($output_bits: tt)*] [$bit: tt $($input_bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name [$bit $($output_bits)*][$($input_bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    (Reverse $name: ident $bits: tt []
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Const $name { len: 0, val: 0 }$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits in reversed order
    // - Output is the constant bit field definition with the sequence of bits as value
    (Const $name: ident { len: $len: expr, val: $val: expr }[]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Const = Const {
                field: Field { pos: $pos, len: $len },
                value: $val,
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    (Const $name: ident { len: $len: expr, val: $val: expr }[$bit: tt $($bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Const $name { len: $len + 1, val: $val * 2 + $bit }[$($bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
}

/// Counts the number of bits equal to zero in a byte slice.
pub fn count_zeros(slice: &[u8]) -> usize {
    slice.iter().map(|&x| x.count_zeros() as usize).sum()
}

/// Returns the number of bits necessary to represent a number.
pub const fn num_bits(x: usize) -> usize {
    8 * core::mem::size_of::<usize>() - x.leading_zeros() as usize
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn field_ok() {
        let field = Field { pos: 3, len: 5 };
        assert_eq!(field.get(0x00000000), 0);
        assert_eq!(field.get(0x00000007), 0);
        assert_eq!(field.get(0x00000008), 1);
        assert_eq!(field.get(0x000000f8), 0x1f);
        assert_eq!(field.get(0x0000ff37), 6);
        let mut word = 0xffffffff;
        field.set(&mut word, 3);
        assert_eq!(word, 0xffffff1f);
    }

    #[test]
    fn const_ok() {
        let field = Const {
            field: Field { pos: 3, len: 5 },
            value: 9,
        };
        assert!(!field.check(0x00000000));
        assert!(!field.check(0x0000ffff));
        assert!(field.check(0x00000048));
        assert!(field.check(0x0000ff4f));
        let mut word = 0xffffffff;
        field.set(&mut word);
        assert_eq!(word, 0xffffff4f);
    }

    #[test]
    fn bit_ok() {
        let bit = Bit { pos: 3 };
        assert!(bit.get(0x00000000));
        assert!(bit.get(0xfffffff7));
        assert!(!bit.get(0x00000008));
        assert!(!bit.get(0xffffffff));
        let mut word = 0xffffffff;
        bit.set(&mut word);
        assert_eq!(word, 0xfffffff7);
    }

    #[test]
    fn checksum_ok() {
        let field = Checksum {
            field: Field { pos: 3, len: 5 },
        };
        assert_eq!(field.get(0x00000000), Err(StoreError::InvalidStorage));
        assert_eq!(field.get(0xffffffff), Ok(31));
        assert_eq!(field.get(0xffffff07), Ok(0));
        assert_eq!(field.get(0xffffff0f), Ok(1));
        assert_eq!(field.get(0x00ffff67), Ok(4));
        assert_eq!(field.get(0x7fffff07), Err(StoreError::InvalidStorage));
        let mut word = 0x0fffffff;
        field.set(&mut word, 4);
        assert_eq!(word, 0x0fffff47);
    }

    #[test]
    fn bitfield_ok() {
        bitfield! {
            FIELD: Field <= 127,
            CONST: Const = [0 1 0 1],
            BIT: Bit,
            CHECKSUM: Checksum <= 58,
            LENGTH: Length,
        }
        assert_eq!(FIELD.pos, 0);
        assert_eq!(FIELD.len, 7);
        assert_eq!(CONST.field.pos, 7);
        assert_eq!(CONST.field.len, 4);
        assert_eq!(CONST.value, 10);
        assert_eq!(BIT.pos, 11);
        assert_eq!(CHECKSUM.field.pos, 12);
        assert_eq!(CHECKSUM.field.len, 6);
        assert_eq!(LENGTH.pos, 18);
    }

    #[test]
    fn count_zeros_ok() {
        assert_eq!(count_zeros(&[0xff, 0xff]), 0);
        assert_eq!(count_zeros(&[0xff, 0xfe]), 1);
        assert_eq!(count_zeros(&[0x7f, 0xff]), 1);
        assert_eq!(count_zeros(&[0x12, 0x48]), 12);
        assert_eq!(count_zeros(&[0x00, 0x00]), 16);
    }

    #[test]
    fn num_bits_ok() {
        assert_eq!(num_bits(0), 0);
        assert_eq!(num_bits(1), 1);
        assert_eq!(num_bits(2), 2);
        assert_eq!(num_bits(3), 2);
        assert_eq!(num_bits(4), 3);
        assert_eq!(num_bits(5), 3);
        assert_eq!(num_bits(8), 4);
        assert_eq!(num_bits(9), 4);
        assert_eq!(num_bits(16), 5);
    }
}
<commit_msg>Clarify documentation<commit_after>// Copyright 2019-2020 Google LLC
//
// 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.

//! Helps manipulate bit fields in 32-bits words.
// TODO(ia0): Remove when the module is used.
#![cfg_attr(not(test), allow(dead_code, unused_macros))]

use crate::{StoreError, StoreResult};

/// Represents a bit field.
///
/// A bit field is a contiguous sequence of bits in a 32-bits word.
///
/// # Invariant
///
/// - The bit field must fit in a 32-bits word: `pos + len < 32`.
pub struct Field {
    /// The position of the bit field.
    pub pos: usize,

    /// The length of the bit field.
    pub len: usize,
}

impl Field {
    /// Reads the value of a bit field.
    pub fn get(&self, word: u32) -> usize {
        ((word >> self.pos) & self.mask()) as usize
    }

    /// Sets the value of a bit field.
    ///
    /// # Preconditions
    ///
    /// - The value must fit in the bit field: `num_bits(value) < self.len`.
    /// - The value must only change bits from 1 to 0: `self.get(*word) & value == value`.
    pub fn set(&self, word: &mut u32, value: usize) {
        let value = value as u32;
        debug_assert_eq!(value & self.mask(), value);
        let mask = !(self.mask() << self.pos);
        *word &= mask | (value << self.pos);
        debug_assert_eq!(self.get(*word), value as usize);
    }

    /// Returns a bit mask the length of the bit field.
    ///
    /// The mask is meant to be applied on a value. It should be shifted to be applied to the bit
    /// field.
    fn mask(&self) -> u32 {
        (1 << self.len) - 1
    }
}

/// Represents a constant bit field.
///
/// # Invariant
///
/// - The value must fit in the bit field: `num_bits(value) <= field.len`.
pub struct Const {
    /// The bit field.
    pub field: Field,

    /// The constant value.
    pub value: usize,
}

impl Const {
    /// Checks that the bit field has its value.
    pub fn check(&self, word: u32) -> bool {
        self.field.get(word) == self.value
    }

    /// Sets the bit field to its value.
    pub fn set(&self, word: &mut u32) {
        self.field.set(word, self.value);
    }
}

/// Represents a single bit.
///
/// # Invariant
///
/// - The bit must fit in a 32-bits word: `pos < 32`.
pub struct Bit {
    /// The position of the bit.
    pub pos: usize,
}

impl Bit {
    /// Returns whether the value of the bit is zero.
    pub fn get(&self, word: u32) -> bool {
        word & (1 << self.pos) == 0
    }

    /// Sets the value of the bit to zero.
    pub fn set(&self, word: &mut u32) {
        *word &= !(1 << self.pos);
    }
}

/// Represents a checksum.
///
/// A checksum is a bit field counting how many bits are set to zero in the word (except in the
/// checksum itself) plus some external increment. It essentially behaves like a bit field storing
/// the external increment.
pub struct Checksum {
    /// The bit field
    pub field: Field,
}

impl Checksum {
    /// Reads the external increment from the checksum.
    ///
    /// # Errors
    ///
    /// Returns `InvalidStorage` if the external increment would be negative.
    pub fn get(&self, word: u32) -> StoreResult<usize> {
        let checksum = self.field.get(word);
        let zeros = word.count_zeros() as usize - (self.field.len - checksum.count_ones() as usize);
        checksum
            .checked_sub(zeros)
            .ok_or(StoreError::InvalidStorage)
    }

    /// Sets the checksum to the external increment value.
    ///
    /// # Preconditions
    ///
    /// - The checksum bits should be set to one.
    /// - The number of zeros in the word plus the value should fit in the checksum.
    pub fn set(&self, word: &mut u32, value: usize) {
        debug_assert_eq!(self.field.get(*word), self.field.mask() as usize);
        self.field.set(word, word.count_zeros() as usize + value);
    }
}

/// Tracks the number of bits used so far.
///
/// # Features
///
/// Only available for tests.
#[cfg(any(doc, test))]
pub struct Length {
    /// The position of the next available bit.
    pub pos: usize,
}

/// Helps defining contiguous bit fields.
///
/// It takes a sequence of bit field descriptors as argument. A bit field descriptor is one of the
/// following:
/// - `$name: Bit,` to define a bit
/// - `$name: Field <= $max,` to define a bit field of minimum length to store `$max`
/// - `$name: Checksum <= $max,` to define a checksum of minimum length to store `$max`
/// - `$name: Length,` to define a length tracker
/// - `$name: Const = [$bits],` to define a constant bit field with value `$bits` (a sequence of
///   space-separated bits)
#[cfg_attr(doc, macro_export)] // For `cargo doc` to produce documentation.
macro_rules! bitfield {
    ($($input: tt)*) => {
        bitfield_impl! { []{ pos: 0 }[$($input)*] }
    };
}

macro_rules! bitfield_impl {
    // Main rules:
    // - Input are bit field descriptors
    // - Position is the number of bits used by prior bit fields
    // - Output are the bit field definitions
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Bit, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Bit = Bit { pos: $pos };]
            { pos: $pos + 1 }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Field <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Field = Field { pos: $pos, len: num_bits($max) };]
            { pos: $pos + $name.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$name: ident: Checksum <= $max: expr, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Checksum = Checksum {
                field: Field { pos: $pos, len: num_bits($max) }
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }
     [$(#[$meta: meta])* $name: ident: Length, $($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* $(#[$meta])* const $name: Length = Length { pos: $pos };]
            { pos: $pos }
            [$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[$name: ident: Const = $bits: tt, $($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name []$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    ([$($output: tt)*]{ pos: $pos: expr }[]) => { $($output)* };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits
    // - Output is the reversed sequence of bits
    (Reverse $name: ident [$($output_bits: tt)*] [$bit: tt $($input_bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Reverse $name [$bit $($output_bits)*][$($input_bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
    (Reverse $name: ident $bits: tt []
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Const $name { len: 0, val: 0 }$bits
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };

    // Auxiliary rules for constant bit fields:
    // - Input is a sequence of bits in reversed order
    // - Output is the constant bit field definition with the sequence of bits as value
    (Const $name: ident { len: $len: expr, val: $val: expr }[]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            [$($output)* const $name: Const = Const {
                field: Field { pos: $pos, len: $len },
                value: $val,
            };]
            { pos: $pos + $name.field.len }
            [$($input)*]
        }
    };
    (Const $name: ident { len: $len: expr, val: $val: expr }[$bit: tt $($bits: tt)*]
     [$($output: tt)*]{ pos: $pos: expr }[$($input: tt)*]) => {
        bitfield_impl! {
            Const $name { len: $len + 1, val: $val * 2 + $bit }[$($bits)*]
            [$($output)*]{ pos: $pos }[$($input)*]
        }
    };
}

/// Counts the number of bits equal to zero in a byte slice.
pub fn count_zeros(slice: &[u8]) -> usize {
    slice.iter().map(|&x| x.count_zeros() as usize).sum()
}

/// Returns the number of bits necessary to represent a number.
pub const fn num_bits(x: usize) -> usize {
    8 * core::mem::size_of::<usize>() - x.leading_zeros() as usize
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn field_ok() {
        let field = Field { pos: 3, len: 5 };
        assert_eq!(field.get(0x00000000), 0);
        assert_eq!(field.get(0x00000007), 0);
        assert_eq!(field.get(0x00000008), 1);
        assert_eq!(field.get(0x000000f8), 0x1f);
        assert_eq!(field.get(0x0000ff37), 6);
        let mut word = 0xffffffff;
        field.set(&mut word, 3);
        assert_eq!(word, 0xffffff1f);
    }

    #[test]
    fn const_ok() {
        let field = Const {
            field: Field { pos: 3, len: 5 },
            value: 9,
        };
        assert!(!field.check(0x00000000));
        assert!(!field.check(0x0000ffff));
        assert!(field.check(0x00000048));
        assert!(field.check(0x0000ff4f));
        let mut word = 0xffffffff;
        field.set(&mut word);
        assert_eq!(word, 0xffffff4f);
    }

    #[test]
    fn bit_ok() {
        let bit = Bit { pos: 3 };
        assert!(bit.get(0x00000000));
        assert!(bit.get(0xfffffff7));
        assert!(!bit.get(0x00000008));
        assert!(!bit.get(0xffffffff));
        let mut word = 0xffffffff;
        bit.set(&mut word);
        assert_eq!(word, 0xfffffff7);
    }

    #[test]
    fn checksum_ok() {
        let field = Checksum {
            field: Field { pos: 3, len: 5 },
        };
        assert_eq!(field.get(0x00000000), Err(StoreError::InvalidStorage));
        assert_eq!(field.get(0xffffffff), Ok(31));
        assert_eq!(field.get(0xffffff07), Ok(0));
        assert_eq!(field.get(0xffffff0f), Ok(1));
        assert_eq!(field.get(0x00ffff67), Ok(4));
        assert_eq!(field.get(0x7fffff07), Err(StoreError::InvalidStorage));
        let mut word = 0x0fffffff;
        field.set(&mut word, 4);
        assert_eq!(word, 0x0fffff47);
    }

    #[test]
    fn bitfield_ok() {
        bitfield! {
            FIELD: Field <= 127,
            CONST: Const = [0 1 0 1],
            BIT: Bit,
            CHECKSUM: Checksum <= 58,
            LENGTH: Length,
        }
        assert_eq!(FIELD.pos, 0);
        assert_eq!(FIELD.len, 7);
        assert_eq!(CONST.field.pos, 7);
        assert_eq!(CONST.field.len, 4);
        assert_eq!(CONST.value, 10);
        assert_eq!(BIT.pos, 11);
        assert_eq!(CHECKSUM.field.pos, 12);
        assert_eq!(CHECKSUM.field.len, 6);
        assert_eq!(LENGTH.pos, 18);
    }

    #[test]
    fn count_zeros_ok() {
        assert_eq!(count_zeros(&[0xff, 0xff]), 0);
        assert_eq!(count_zeros(&[0xff, 0xfe]), 1);
        assert_eq!(count_zeros(&[0x7f, 0xff]), 1);
        assert_eq!(count_zeros(&[0x12, 0x48]), 12);
        assert_eq!(count_zeros(&[0x00, 0x00]), 16);
    }

    #[test]
    fn num_bits_ok() {
        assert_eq!(num_bits(0), 0);
        assert_eq!(num_bits(1), 1);
        assert_eq!(num_bits(2), 2);
        assert_eq!(num_bits(3), 2);
        assert_eq!(num_bits(4), 3);
        assert_eq!(num_bits(5), 3);
        assert_eq!(num_bits(8), 4);
        assert_eq!(num_bits(9), 4);
        assert_eq!(num_bits(16), 5);
    }
}
<|endoftext|>


<commit_before>use super::{child_properties, function, general, properties, signal, trait_impls};
use crate::{
    analysis::{
        self,
        rust_type::{rust_type, rust_type_full},
    },
    case::CaseExt,
    codegen::general::version_condition_string,
    env::Env,
    library, nameutil,
    traits::IntoString,
};
use std::io::{Result, Write};

pub fn generate(
    w: &mut dyn Write,
    env: &Env,
    analysis: &analysis::object::Info,
    generate_display_trait: bool,
) -> Result<()> {
    general::start_comments(w, &env.config)?;
    general::uses(w, env, &analysis.imports)?;

    general::define_object_type(
        w,
        env,
        &analysis.name,
        &analysis.c_type,
        &analysis.c_class_type.as_ref().map(|s| &s[..]),
        &analysis.rust_class_type.as_ref().map(|s| &s[..]),
        &analysis.get_type,
        analysis.is_interface,
        &analysis.supertypes,
    )?;

    if need_generate_inherent(analysis) {
        writeln!(w)?;
        write!(w, "impl {} {{", analysis.name)?;
        for func_analysis in &analysis.constructors() {
            function::generate(w, env, func_analysis, false, false, 1)?;
        }

        if !need_generate_trait(analysis) {
            for func_analysis in &analysis.methods() {
                function::generate(w, env, func_analysis, false, false, 1)?;
            }

            for property in &analysis.properties {
                properties::generate(w, env, property, false, false, 1)?;
            }

            for child_property in &analysis.child_properties {
                child_properties::generate(w, env, child_property, false, false, 1)?;
            }
        }

        for func_analysis in &analysis.functions() {
            function::generate(w, env, func_analysis, false, false, 1)?;
        }

        if !need_generate_trait(analysis) {
            for signal_analysis in analysis
                .signals
                .iter()
                .chain(analysis.notify_signals.iter())
            {
                signal::generate(w, env, signal_analysis, false, false, 1)?;
            }
        }

        writeln!(w, "}}")?;

        general::declare_default_from_new(w, env, &analysis.name, &analysis.functions)?;
    }

    trait_impls::generate(
        w,
        &analysis.name,
        &analysis.functions,
        &analysis.specials,
        if need_generate_trait(analysis) {
            Some(&analysis.trait_name)
        } else {
            None
        },
    )?;

    if !analysis.builder_properties.is_empty() {
        writeln!(w)?;
        generate_builder(w, env, analysis)?;
    }

    if analysis.concurrency != library::Concurrency::None {
        writeln!(w)?;
    }

    match analysis.concurrency {
        library::Concurrency::Send | library::Concurrency::SendSync => {
            writeln!(w, "unsafe impl Send for {} {{}}", analysis.name)?;
        }
        library::Concurrency::SendUnique => {
            if env.namespaces.is_glib_crate {
                writeln!(w, "unsafe impl ::SendUnique for {} {{", analysis.name)?;
            } else {
                writeln!(w, "unsafe impl glib::SendUnique for {} {{", analysis.name)?;
            }

            writeln!(w, "    fn is_unique(&self) -> bool {{")?;
            writeln!(w, "        self.ref_count() == 1")?;
            writeln!(w, "    }}")?;

            writeln!(w, "}}")?;
        }
        _ => (),
    }

    if let library::Concurrency::SendSync = analysis.concurrency {
        writeln!(w, "unsafe impl Sync for {} {{}}", analysis.name)?;
    }

    if !analysis.final_type {
        writeln!(w)?;
        writeln!(
            w,
            "pub const NONE_{}: Option<&{}> = None;",
            analysis.name.to_snake().to_uppercase(),
            analysis.name
        )?;
    }

    if need_generate_trait(analysis) {
        writeln!(w)?;
        generate_trait(w, env, analysis)?;
    }

    if generate_display_trait {
        writeln!(w, "\nimpl fmt::Display for {} {{", analysis.name,)?;
        // Generate Display trait implementation.
        writeln!(
            w,
            "\tfn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{\n\
             \t\twrite!(f, \"{}\")\n\
             \t}}\n\
             }}",
            analysis.name
        )?;
    }

    Ok(())
}

// TODO: instead create a Vec<> inside the Builder instead of Options.
fn generate_builder(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    let mut methods = vec![];
    let mut properties = vec![];
    writeln!(w, "#[derive(Clone, Default)]")?;
    writeln!(w, "pub struct {}Builder {{", analysis.name)?;
    for property in &analysis.builder_properties {
        match rust_type(env, property.typ) {
            Ok(type_string) => {
                let type_string = match type_string.as_str() {
                    "GString" => "String",
                    "Vec<GString>" => "Vec<String>",
                    typ => typ,
                };
                let mut param_type = rust_type_full(
                    env,
                    property.typ,
                    library::Nullable(false),
                    property.set_in_ref_mode,
                    library::ParameterScope::None,
                    library::Concurrency::None,
                )
                .into_string();
                let (param_type_override, bounds, conversion) = match &param_type[..] {
                    "&str" => (None, String::new(), ".to_string()"),
                    "&[&str]" => (Some("Vec<String>".to_string()), String::new(), ""),
                    _ if !property.bounds.is_empty() => {
                        let (bounds, _) = function::bounds(&property.bounds, &[], false, false);
                        let alias = property
                            .bounds
                            .get_parameter_alias_info(&property.name)
                            .map(|(alias, _)| format!("&{}", alias));
                        (alias, bounds, ".clone().upcast()")
                    }
                    typ if typ.starts_with('&') => (None, String::new(), ".clone()"),
                    _ => (None, String::new(), ""),
                };
                if let Some(param_type_override) = param_type_override {
                    param_type = param_type_override.to_string();
                }
                let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
                let version_condition_string =
                    version_condition_string(env, property.version, false, 1);
                if let Some(ref version_condition_string) = version_condition_string {
                    writeln!(w, "{}", version_condition_string)?;
                }
                writeln!(w, "    {}: Option<{}>,", name, type_string)?;
                let prefix = version_condition_string
                    .map(|version| format!("{}\n", version))
                    .unwrap_or_default();
                methods.push(format!(
                    "\n{prefix}    pub fn {name}{bounds}(mut self, {name}: {param_type}) -> Self {{
        self.{name} = Some({name}{conversion});
        self
    }}",
                    prefix = prefix,
                    param_type = param_type,
                    name = name,
                    conversion = conversion,
                    bounds = bounds
                ));
                properties.push(property);
            }
            Err(_) => writeln!(w, "    //{}: /*Unknown type*/,", property.name)?,
        }
    }
    writeln!(
        w,
        "}}

impl {}Builder {{
    pub fn new() -> Self {{
        Self::default()
    }}
",
        analysis.name
    )?;

    writeln!(
        w,
        "
    pub fn build(self) -> {} {{
        let mut properties: Vec<(&str, &dyn ToValue)> = vec![];",
        analysis.name
    )?;
    for property in &properties {
        let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
        let version_condition_string = version_condition_string(env, property.version, false, 2);
        let condition_tabs = if version_condition_string.is_some() {
            "\t"
        } else {
            ""
        };
        if let Some(ref version_condition_string) = version_condition_string {
            writeln!(w, "{}", version_condition_string)?;
            writeln!(w, "        {{")?;
        }
        writeln!(
            w,
            "{tabs}        if let Some(ref {field}) = self.{field} {{
{tabs}            properties.push((\"{name}\", {field}));
{tabs}        }}",
            name = property.name,
            field = name,
            tabs = condition_tabs
        )?;
        if version_condition_string.is_some() {
            writeln!(w, "        }}")?;
        }
    }
    let glib_crate_name = if env.namespaces.is_glib_crate {
        "crate"
    } else {
        "glib"
    };
    writeln!(w,
"        {}::Object::new({}::static_type(), &properties).expect(\"object new\").downcast().expect(\"downcast\")
    }}", glib_crate_name, analysis.name)?;
    for method in methods {
        writeln!(w, "{}", method)?;
    }
    writeln!(w, "}}")
}

fn generate_trait(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    write!(w, "pub trait {}: 'static {{", analysis.trait_name)?;

    for func_analysis in &analysis.methods() {
        function::generate(w, env, func_analysis, true, true, 1)?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, true, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, true, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, true, 1)?;
    }
    writeln!(w, "}}")?;

    writeln!(w)?;
    write!(
        w,
        "impl<O: IsA<{}>> {} for O {{",
        analysis.name, analysis.trait_name,
    )?;

    for func_analysis in &analysis.methods() {
        function::generate(w, env, func_analysis, true, false, 1)?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, false, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, false, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, false, 1)?;
    }
    writeln!(w, "}}")?;

    Ok(())
}

fn need_generate_inherent(analysis: &analysis::object::Info) -> bool {
    analysis.has_constructors || analysis.has_functions || !need_generate_trait(analysis)
}

fn need_generate_trait(analysis: &analysis::object::Info) -> bool {
    analysis.generate_trait
}

pub fn generate_reexports(
    env: &Env,
    analysis: &analysis::object::Info,
    module_name: &str,
    contents: &mut Vec<String>,
    traits: &mut Vec<String>,
) {
    let mut cfgs: Vec<String> = Vec::new();
    if let Some(cfg) = general::cfg_condition_string(&analysis.cfg_condition, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::version_condition_string(env, analysis.version, false, 0) {
        cfgs.push(cfg);
    }
    contents.push("".to_owned());
    contents.extend_from_slice(&cfgs);
    contents.push(format!("mod {};", module_name));
    contents.extend_from_slice(&cfgs);

    let none_type = if !analysis.final_type {
        format!(", NONE_{}", analysis.name.to_snake().to_uppercase())
    } else {
        String::new()
    };

    if let Some(ref class_name) = analysis.rust_class_type {
        contents.push(format!(
            "pub use self::{}::{{{}, {}{}}};",
            module_name, analysis.name, class_name, none_type
        ));
    } else {
        contents.push(format!(
            "pub use self::{}::{{{}{}}};",
            module_name, analysis.name, none_type
        ));
    }
    if need_generate_trait(analysis) {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{};",
            module_name, analysis.trait_name
        ));
        for cfg in &cfgs {
            traits.push(format!("\t{}", cfg));
        }
        traits.push(format!("\tpub use super::{};", analysis.trait_name));
    }

    if !analysis.builder_properties.is_empty() {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{}Builder;",
            module_name, analysis.name
        ));
    }
}
<commit_msg>codegen: Use new builder_postprocess attribute<commit_after>use super::{child_properties, function, general, properties, signal, trait_impls};
use crate::{
    analysis::{
        self,
        rust_type::{rust_type, rust_type_full},
    },
    case::CaseExt,
    codegen::general::version_condition_string,
    env::Env,
    library, nameutil,
    traits::IntoString,
};
use std::io::{Result, Write};

pub fn generate(
    w: &mut dyn Write,
    env: &Env,
    analysis: &analysis::object::Info,
    generate_display_trait: bool,
) -> Result<()> {
    general::start_comments(w, &env.config)?;
    general::uses(w, env, &analysis.imports)?;

    general::define_object_type(
        w,
        env,
        &analysis.name,
        &analysis.c_type,
        &analysis.c_class_type.as_ref().map(|s| &s[..]),
        &analysis.rust_class_type.as_ref().map(|s| &s[..]),
        &analysis.get_type,
        analysis.is_interface,
        &analysis.supertypes,
    )?;

    if need_generate_inherent(analysis) {
        writeln!(w)?;
        write!(w, "impl {} {{", analysis.name)?;
        for func_analysis in &analysis.constructors() {
            function::generate(w, env, func_analysis, false, false, 1)?;
        }

        if !need_generate_trait(analysis) {
            for func_analysis in &analysis.methods() {
                function::generate(w, env, func_analysis, false, false, 1)?;
            }

            for property in &analysis.properties {
                properties::generate(w, env, property, false, false, 1)?;
            }

            for child_property in &analysis.child_properties {
                child_properties::generate(w, env, child_property, false, false, 1)?;
            }
        }

        for func_analysis in &analysis.functions() {
            function::generate(w, env, func_analysis, false, false, 1)?;
        }

        if !need_generate_trait(analysis) {
            for signal_analysis in analysis
                .signals
                .iter()
                .chain(analysis.notify_signals.iter())
            {
                signal::generate(w, env, signal_analysis, false, false, 1)?;
            }
        }

        writeln!(w, "}}")?;

        general::declare_default_from_new(w, env, &analysis.name, &analysis.functions)?;
    }

    trait_impls::generate(
        w,
        &analysis.name,
        &analysis.functions,
        &analysis.specials,
        if need_generate_trait(analysis) {
            Some(&analysis.trait_name)
        } else {
            None
        },
    )?;

    if !analysis.builder_properties.is_empty() {
        writeln!(w)?;
        generate_builder(w, env, analysis)?;
    }

    if analysis.concurrency != library::Concurrency::None {
        writeln!(w)?;
    }

    match analysis.concurrency {
        library::Concurrency::Send | library::Concurrency::SendSync => {
            writeln!(w, "unsafe impl Send for {} {{}}", analysis.name)?;
        }
        library::Concurrency::SendUnique => {
            if env.namespaces.is_glib_crate {
                writeln!(w, "unsafe impl ::SendUnique for {} {{", analysis.name)?;
            } else {
                writeln!(w, "unsafe impl glib::SendUnique for {} {{", analysis.name)?;
            }

            writeln!(w, "    fn is_unique(&self) -> bool {{")?;
            writeln!(w, "        self.ref_count() == 1")?;
            writeln!(w, "    }}")?;

            writeln!(w, "}}")?;
        }
        _ => (),
    }

    if let library::Concurrency::SendSync = analysis.concurrency {
        writeln!(w, "unsafe impl Sync for {} {{}}", analysis.name)?;
    }

    if !analysis.final_type {
        writeln!(w)?;
        writeln!(
            w,
            "pub const NONE_{}: Option<&{}> = None;",
            analysis.name.to_snake().to_uppercase(),
            analysis.name
        )?;
    }

    if need_generate_trait(analysis) {
        writeln!(w)?;
        generate_trait(w, env, analysis)?;
    }

    if generate_display_trait {
        writeln!(w, "\nimpl fmt::Display for {} {{", analysis.name,)?;
        // Generate Display trait implementation.
        writeln!(
            w,
            "\tfn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{\n\
             \t\twrite!(f, \"{}\")\n\
             \t}}\n\
             }}",
            analysis.name
        )?;
    }

    Ok(())
}

// TODO: instead create a Vec<> inside the Builder instead of Options.
fn generate_builder(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    let mut methods = vec![];
    let mut properties = vec![];
    writeln!(w, "#[derive(Clone, Default)]")?;
    writeln!(w, "pub struct {}Builder {{", analysis.name)?;
    for property in &analysis.builder_properties {
        match rust_type(env, property.typ) {
            Ok(type_string) => {
                let type_string = match type_string.as_str() {
                    "GString" => "String",
                    "Vec<GString>" => "Vec<String>",
                    typ => typ,
                };
                let mut param_type = rust_type_full(
                    env,
                    property.typ,
                    library::Nullable(false),
                    property.set_in_ref_mode,
                    library::ParameterScope::None,
                    library::Concurrency::None,
                )
                .into_string();
                let (param_type_override, bounds, conversion) = match &param_type[..] {
                    "&str" => (None, String::new(), ".to_string()"),
                    "&[&str]" => (Some("Vec<String>".to_string()), String::new(), ""),
                    _ if !property.bounds.is_empty() => {
                        let (bounds, _) = function::bounds(&property.bounds, &[], false, false);
                        let alias = property
                            .bounds
                            .get_parameter_alias_info(&property.name)
                            .map(|(alias, _)| format!("&{}", alias));
                        (alias, bounds, ".clone().upcast()")
                    }
                    typ if typ.starts_with('&') => (None, String::new(), ".clone()"),
                    _ => (None, String::new(), ""),
                };
                if let Some(param_type_override) = param_type_override {
                    param_type = param_type_override.to_string();
                }
                let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
                let version_condition_string =
                    version_condition_string(env, property.version, false, 1);
                if let Some(ref version_condition_string) = version_condition_string {
                    writeln!(w, "{}", version_condition_string)?;
                }
                writeln!(w, "    {}: Option<{}>,", name, type_string)?;
                let prefix = version_condition_string
                    .map(|version| format!("{}\n", version))
                    .unwrap_or_default();
                methods.push(format!(
                    "\n{prefix}    pub fn {name}{bounds}(mut self, {name}: {param_type}) -> Self {{
        self.{name} = Some({name}{conversion});
        self
    }}",
                    prefix = prefix,
                    param_type = param_type,
                    name = name,
                    conversion = conversion,
                    bounds = bounds
                ));
                properties.push(property);
            }
            Err(_) => writeln!(w, "    //{}: /*Unknown type*/,", property.name)?,
        }
    }
    writeln!(
        w,
        "}}

impl {}Builder {{
    pub fn new() -> Self {{
        Self::default()
    }}
",
        analysis.name
    )?;

    writeln!(
        w,
        "
    pub fn build(self) -> {} {{
        let mut properties: Vec<(&str, &dyn ToValue)> = vec![];",
        analysis.name
    )?;
    for property in &properties {
        let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
        let version_condition_string = version_condition_string(env, property.version, false, 2);
        let condition_tabs = if version_condition_string.is_some() {
            "\t"
        } else {
            ""
        };
        if let Some(ref version_condition_string) = version_condition_string {
            writeln!(w, "{}", version_condition_string)?;
            writeln!(w, "        {{")?;
        }
        writeln!(
            w,
            "{tabs}        if let Some(ref {field}) = self.{field} {{
{tabs}            properties.push((\"{name}\", {field}));
{tabs}        }}",
            name = property.name,
            field = name,
            tabs = condition_tabs
        )?;
        if version_condition_string.is_some() {
            writeln!(w, "        }}")?;
        }
    }
    let glib_crate_name = if env.namespaces.is_glib_crate {
        "crate"
    } else {
        "glib"
    };
    writeln!(
        w,
        r#"        let ret = {}::Object::new({}::static_type(), &properties)
            .expect("object new")
            .downcast::<{}>()
            .expect("downcast");"#,
        glib_crate_name, analysis.name, analysis.name
    )?;
    if let Some(code) = analysis.builder_postprocess.as_ref() {
        writeln!(w, "        {{\n            {}\n        }}", code)?;
    }
    writeln!(w, "    ret\n    }}")?;
    for method in methods {
        writeln!(w, "{}", method)?;
    }
    writeln!(w, "}}")
}

fn generate_trait(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    write!(w, "pub trait {}: 'static {{", analysis.trait_name)?;

    for func_analysis in &analysis.methods() {
        function::generate(w, env, func_analysis, true, true, 1)?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, true, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, true, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, true, 1)?;
    }
    writeln!(w, "}}")?;

    writeln!(w)?;
    write!(
        w,
        "impl<O: IsA<{}>> {} for O {{",
        analysis.name, analysis.trait_name,
    )?;

    for func_analysis in &analysis.methods() {
        function::generate(w, env, func_analysis, true, false, 1)?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, false, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, false, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, false, 1)?;
    }
    writeln!(w, "}}")?;

    Ok(())
}

fn need_generate_inherent(analysis: &analysis::object::Info) -> bool {
    analysis.has_constructors || analysis.has_functions || !need_generate_trait(analysis)
}

fn need_generate_trait(analysis: &analysis::object::Info) -> bool {
    analysis.generate_trait
}

pub fn generate_reexports(
    env: &Env,
    analysis: &analysis::object::Info,
    module_name: &str,
    contents: &mut Vec<String>,
    traits: &mut Vec<String>,
) {
    let mut cfgs: Vec<String> = Vec::new();
    if let Some(cfg) = general::cfg_condition_string(&analysis.cfg_condition, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::version_condition_string(env, analysis.version, false, 0) {
        cfgs.push(cfg);
    }
    contents.push("".to_owned());
    contents.extend_from_slice(&cfgs);
    contents.push(format!("mod {};", module_name));
    contents.extend_from_slice(&cfgs);

    let none_type = if !analysis.final_type {
        format!(", NONE_{}", analysis.name.to_snake().to_uppercase())
    } else {
        String::new()
    };

    if let Some(ref class_name) = analysis.rust_class_type {
        contents.push(format!(
            "pub use self::{}::{{{}, {}{}}};",
            module_name, analysis.name, class_name, none_type
        ));
    } else {
        contents.push(format!(
            "pub use self::{}::{{{}{}}};",
            module_name, analysis.name, none_type
        ));
    }
    if need_generate_trait(analysis) {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{};",
            module_name, analysis.trait_name
        ));
        for cfg in &cfgs {
            traits.push(format!("\t{}", cfg));
        }
        traits.push(format!("\tpub use super::{};", analysis.trait_name));
    }

    if !analysis.builder_properties.is_empty() {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{}Builder;",
            module_name, analysis.name
        ));
    }
}
<|endoftext|>