<commit_before>use types::*;
use types::value::*;
use error::{Result as CResult, Error as CError};
use consistency::Consistency;
use frame::frame_query::{ParamsReqQuery, QueryFlags};
use frame::frame_batch::{BatchType, BatchQuery, BodyReqBatch, BatchQuerySubj};

/// instead of writing functions which resemble
/// ```
/// pub fn query<'a> (&'a mut self,query: String) -> &'a mut Self{
///     self.query = Some(query);
///            self
/// }
/// ```
/// and repeating it for all the attributes; it is extracted out as a macro so that code
/// is more concise see
/// @https://doc.rust-lang.org/book/method-syntax.html
///
///
///
macro_rules! builder_opt_field {
    ($field:ident, $field_type:ty) => {
        pub fn $field(mut self,
                          $field: $field_type) -> Self {
            self.$field = Some($field);
            self
        }
    };
}

/// Structure that represents CQL query and parameters which will be applied during
/// its execution
#[derive(Debug, Default)]
pub struct Query {
    pub query: String,
    // query parameters
    pub consistency: Consistency,
    pub values: Option<Vec<Value>>,
    pub with_names: Option<bool>,
    pub page_size: Option<i32>,
    pub paging_state: Option<CBytes>,
    pub serial_consistency: Option<Consistency>,
    pub timestamp: Option<i32>,
}

/// QueryBuilder is a helper sturcture that helps to construct `Query`. `Query` itself
/// consists of CQL query string and list of parameters.
/// Parameters are the same as ones described in [Cassandra v4 protocol]
/// (https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v4.spec#L304)
#[derive(Debug, Default)]
pub struct QueryBuilder {
    query: String,
    consistency: Consistency,
    values: Option<Vec<Value>>,
    with_names: Option<bool>,
    page_size: Option<i32>,
    paging_state: Option<CBytes>,
    serial_consistency: Option<Consistency>,
    timestamp: Option<i32>,
}

impl QueryBuilder {
    /// Factory function that takes CQL as an argument and returns new `QueryBuilder`.
    /// Default consistency level is `One`
    pub fn new<T: Into<String>>(query: T) -> QueryBuilder {
        return QueryBuilder { query: query.into(), ..Default::default() };
    }

    /// Sets new query consistency
    pub fn consistency(mut self, consistency: Consistency) -> Self {
        self.consistency = consistency;

        return self;
    }

    /// Sets new query values
    builder_opt_field!(values, Vec<Value>);

    /// Sets new query with_names
    builder_opt_field!(with_names, bool);

    /// Sets new query pagesize
    builder_opt_field!(page_size, i32);

    /// Sets new query pagin state
    builder_opt_field!(paging_state, CBytes);

    /// Sets new query serial_consistency
    builder_opt_field!(serial_consistency, Consistency);

    /// Sets new quey timestamp
    builder_opt_field!(timestamp, i32);

    pub fn apply_query_params(mut self, params: QueryParams) -> Self {
        self.consistency = params.consistency;
        self.values = params.values;
        self.page_size = params.page_size;
        self.paging_state = params.paging_state;
        self.serial_consistency = params.serial_consistency;
        self.timestamp = params.timestamp;

        return self;
    }

    /// Finalizes query building process and returns query itself
    pub fn finalize(self) -> Query {
        return Query {
            query: self.query,
            consistency: self.consistency,
            values: self.values,
            with_names: self.with_names,
            page_size: self.page_size,
            paging_state: self.paging_state,
            serial_consistency: self.serial_consistency,
            timestamp: self.timestamp,
        };
    }
}

pub type QueryParams = ParamsReqQuery;

/// Query parameters builder
#[derive(Debug)]
pub struct QueryParamsBuilder {
    consistency: Consistency,
    values: Option<Vec<Value>>,
    with_names: bool,
    page_size: Option<i32>,
    paging_state: Option<CBytes>,
    serial_consistency: Option<Consistency>,
    timestamp: Option<i32>,
}

impl QueryParamsBuilder {
    pub fn new(consistency: Consistency) -> QueryParamsBuilder {
        return QueryParamsBuilder {
            consistency: consistency,
            values: None,
            with_names: false,
            page_size: None,
            paging_state: None,
            serial_consistency: None,
            timestamp: None,
        };
    }

    pub fn values(mut self, v: Vec<Value>) -> Self {
        self.values = Some(v);

        return self;
    }

    pub fn with_names(mut self, with_names: bool) -> Self {
        self.with_names = with_names;

        return self;
    }

    pub fn page_size(mut self, page_size: i32) -> Self {
        self.page_size = Some(page_size);

        return self;
    }

    pub fn paging_state(mut self, paging_state: CBytes) -> Self {
        self.paging_state = Some(paging_state);

        return self;
    }

    pub fn serial_consistency(mut self, serial_consistency: Consistency) -> Self {
        self.serial_consistency = Some(serial_consistency);

        return self;
    }

    pub fn timestamp(mut self, timestamp: i32) -> Self {
        self.timestamp = Some(timestamp);

        return self;
    }

    pub fn finalize(self) -> QueryParams {
        // query flags
        let mut flags: Vec<QueryFlags> = vec![];

        if self.values.is_some() {
            flags.push(QueryFlags::Value);
        }

        if self.with_names {
            flags.push(QueryFlags::WithNamesForValues);
        }

        if self.page_size.is_some() {
            flags.push(QueryFlags::PageSize);
        }

        if self.serial_consistency.is_some() {
            flags.push(QueryFlags::WithSerialConsistency);
        }

        if self.timestamp.is_some() {
            flags.push(QueryFlags::WithDefaultTimestamp);
        }

        //TODO need to revisit// borrow checker was complaining do without the clone
        QueryParams {
            consistency: self.consistency,
            flags: flags,
            values: self.values,
            page_size: self.page_size,
            paging_state: self.paging_state,
            serial_consistency: self.serial_consistency,
            timestamp: self.timestamp,
        }

    }
}

pub type QueryBatch = BodyReqBatch;

#[derive(Debug)]
pub struct BatchQueryBuilder {
    batch_type: BatchType,
    queries: Vec<BatchQuery>,
    consistency: Consistency,
    serial_consistency: Option<Consistency>,
    timestamp: Option<i32>,
}

impl BatchQueryBuilder {
    pub fn new() -> BatchQueryBuilder {
        BatchQueryBuilder {
            batch_type: BatchType::Logged,
            queries: vec![],
            consistency: Consistency::One,
            serial_consistency: None,
            timestamp: None,
        }
    }

    pub fn batch_type(mut self, batch_type: BatchType) -> Self {
        self.batch_type = batch_type;
        self
    }

    /// Add a query (non-prepared one)
    pub fn add_query<T: Into<String>>(mut self, query: T, values: Vec<BatchValue>) -> Self {
        self.queries.push(BatchQuery {
            is_prepared: false,
            subject: BatchQuerySubj::QueryString(CStringLong::new(query.into())),
            values: values,
        });
        self
    }

    /// Add a query (prepared one)
    pub fn add_query_prepared(mut self, query_id: CBytesShort, values: Vec<BatchValue>) -> Self {
        self.queries.push(BatchQuery {
            is_prepared: true,
            subject: BatchQuerySubj::PreparedId(query_id),
            values: values,
        });
        self
    }

    pub fn clear_queries(mut self) -> Self {
        self.queries = vec![];
        self
    }

    pub fn consistency(mut self, consistency: Consistency) -> Self {
        self.consistency = consistency;
        self
    }

    pub fn serial_consistency(mut self, serial_consistency: Option<Consistency>) -> Self {
        self.serial_consistency = serial_consistency;
        self
    }

    pub fn timestamp(mut self, timestamp: Option<i32>) -> Self {
        self.timestamp = timestamp;
        self
    }

    pub fn finalize(self) -> CResult<BodyReqBatch> {
        let mut flags = vec![];

        if self.serial_consistency.is_some() {
            flags.push(QueryFlags::WithSerialConsistency);
        }

        if self.timestamp.is_some() {
            flags.push(QueryFlags::WithDefaultTimestamp);
        }

        let with_names_for_values = self.queries
            .iter()
            .all(|q| q.values.iter().all(|v| v.0.is_some()));

        if !with_names_for_values {
            let some_names_for_values = self.queries
                .iter()
                .any(|q| q.values.iter().any(|v| v.0.is_some()));

            if some_names_for_values {
                return Err(CError::General(String::from("Inconsistent query values - mixed \
                                                         with and without names values")));
            }
        }

        if with_names_for_values {
            flags.push(QueryFlags::WithNamesForValues);
        }

        Ok(BodyReqBatch {
            batch_type: self.batch_type,
            queries: self.queries,
            query_flags: flags,
            consistency: self.consistency,
            serial_consistency: self.serial_consistency,
            timestamp: self.timestamp,
        })
    }
}

pub type BatchValue = (Option<CString>, Value);

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

    #[test]
    fn new() {
        let _ = QueryBuilder::new("USE keyspace").finalize();
    }

    #[test]
    fn with_parameters() {
        let _ = QueryBuilder::new("USE keyspace")
            .consistency(Consistency::Two)
            .values(vec![Value::new_null()])
            .with_names(true)
            .page_size(11)
            .paging_state(CBytes::new(vec![1, 2, 3, 4, 5]))
            .serial_consistency(Consistency::One)
            .timestamp(1245678)
            .finalize();
    }
}

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

    #[test]
    fn new() {
        let _ = QueryParamsBuilder::new(Consistency::Two).finalize();
    }

    #[test]
    fn with_parameters() {
        let _ = QueryParamsBuilder::new(Consistency::Two)
            .values(vec![Value::new_null()])
            .with_names(true)
            .page_size(11)
            .paging_state(CBytes::new(vec![1, 2, 3, 4, 5]))
            .serial_consistency(Consistency::One)
            .timestamp(1245678)
            .finalize();
    }
}

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

    #[test]
    fn new() {
        assert!(BatchQueryBuilder::new().finalize().is_ok());
    }

    #[test]
    fn with_parameters() {
        let q = BatchQueryBuilder::new()
            .batch_type(BatchType::Logged)
            .add_query("some query".to_string(), vec![])
            .add_query_prepared(CBytesShort::new(vec![1, 2, 3]), vec![])
            .consistency(Consistency::One)
            .serial_consistency(Some(Consistency::Two))
            .timestamp(Some(1234))
            .finalize();
        assert!(q.is_ok())
    }

    #[test]
    fn clear_queries() {
        let q = BatchQueryBuilder::new()
            .add_query("some query".to_string(), vec![])
            .clear_queries()
            .finalize();
        assert!(q.is_ok())
    }
}
<commit_msg>Remove outdated todo<commit_after>use types::*;
use types::value::*;
use error::{Result as CResult, Error as CError};
use consistency::Consistency;
use frame::frame_query::{ParamsReqQuery, QueryFlags};
use frame::frame_batch::{BatchType, BatchQuery, BodyReqBatch, BatchQuerySubj};

/// instead of writing functions which resemble
/// ```
/// pub fn query<'a> (&'a mut self,query: String) -> &'a mut Self{
///     self.query = Some(query);
///            self
/// }
/// ```
/// and repeating it for all the attributes; it is extracted out as a macro so that code
/// is more concise see
/// @https://doc.rust-lang.org/book/method-syntax.html
///
///
///
macro_rules! builder_opt_field {
    ($field:ident, $field_type:ty) => {
        pub fn $field(mut self,
                          $field: $field_type) -> Self {
            self.$field = Some($field);
            self
        }
    };
}

/// Structure that represents CQL query and parameters which will be applied during
/// its execution
#[derive(Debug, Default)]
pub struct Query {
    pub query: String,
    // query parameters
    pub consistency: Consistency,
    pub values: Option<Vec<Value>>,
    pub with_names: Option<bool>,
    pub page_size: Option<i32>,
    pub paging_state: Option<CBytes>,
    pub serial_consistency: Option<Consistency>,
    pub timestamp: Option<i32>,
}

/// QueryBuilder is a helper sturcture that helps to construct `Query`. `Query` itself
/// consists of CQL query string and list of parameters.
/// Parameters are the same as ones described in [Cassandra v4 protocol]
/// (https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v4.spec#L304)
#[derive(Debug, Default)]
pub struct QueryBuilder {
    query: String,
    consistency: Consistency,
    values: Option<Vec<Value>>,
    with_names: Option<bool>,
    page_size: Option<i32>,
    paging_state: Option<CBytes>,
    serial_consistency: Option<Consistency>,
    timestamp: Option<i32>,
}

impl QueryBuilder {
    /// Factory function that takes CQL as an argument and returns new `QueryBuilder`.
    /// Default consistency level is `One`
    pub fn new<T: Into<String>>(query: T) -> QueryBuilder {
        return QueryBuilder { query: query.into(), ..Default::default() };
    }

    /// Sets new query consistency
    pub fn consistency(mut self, consistency: Consistency) -> Self {
        self.consistency = consistency;

        return self;
    }

    /// Sets new query values
    builder_opt_field!(values, Vec<Value>);

    /// Sets new query with_names
    builder_opt_field!(with_names, bool);

    /// Sets new query pagesize
    builder_opt_field!(page_size, i32);

    /// Sets new query pagin state
    builder_opt_field!(paging_state, CBytes);

    /// Sets new query serial_consistency
    builder_opt_field!(serial_consistency, Consistency);

    /// Sets new quey timestamp
    builder_opt_field!(timestamp, i32);

    pub fn apply_query_params(mut self, params: QueryParams) -> Self {
        self.consistency = params.consistency;
        self.values = params.values;
        self.page_size = params.page_size;
        self.paging_state = params.paging_state;
        self.serial_consistency = params.serial_consistency;
        self.timestamp = params.timestamp;

        return self;
    }

    /// Finalizes query building process and returns query itself
    pub fn finalize(self) -> Query {
        return Query {
            query: self.query,
            consistency: self.consistency,
            values: self.values,
            with_names: self.with_names,
            page_size: self.page_size,
            paging_state: self.paging_state,
            serial_consistency: self.serial_consistency,
            timestamp: self.timestamp,
        };
    }
}

pub type QueryParams = ParamsReqQuery;

/// Query parameters builder
#[derive(Debug)]
pub struct QueryParamsBuilder {
    consistency: Consistency,
    values: Option<Vec<Value>>,
    with_names: bool,
    page_size: Option<i32>,
    paging_state: Option<CBytes>,
    serial_consistency: Option<Consistency>,
    timestamp: Option<i32>,
}

impl QueryParamsBuilder {
    pub fn new(consistency: Consistency) -> QueryParamsBuilder {
        return QueryParamsBuilder {
            consistency: consistency,
            values: None,
            with_names: false,
            page_size: None,
            paging_state: None,
            serial_consistency: None,
            timestamp: None,
        };
    }

    pub fn values(mut self, v: Vec<Value>) -> Self {
        self.values = Some(v);

        return self;
    }

    pub fn with_names(mut self, with_names: bool) -> Self {
        self.with_names = with_names;

        return self;
    }

    pub fn page_size(mut self, page_size: i32) -> Self {
        self.page_size = Some(page_size);

        return self;
    }

    pub fn paging_state(mut self, paging_state: CBytes) -> Self {
        self.paging_state = Some(paging_state);

        return self;
    }

    pub fn serial_consistency(mut self, serial_consistency: Consistency) -> Self {
        self.serial_consistency = Some(serial_consistency);

        return self;
    }

    pub fn timestamp(mut self, timestamp: i32) -> Self {
        self.timestamp = Some(timestamp);

        return self;
    }

    pub fn finalize(self) -> QueryParams {
        // query flags
        let mut flags: Vec<QueryFlags> = vec![];

        if self.values.is_some() {
            flags.push(QueryFlags::Value);
        }

        if self.with_names {
            flags.push(QueryFlags::WithNamesForValues);
        }

        if self.page_size.is_some() {
            flags.push(QueryFlags::PageSize);
        }

        if self.serial_consistency.is_some() {
            flags.push(QueryFlags::WithSerialConsistency);
        }

        if self.timestamp.is_some() {
            flags.push(QueryFlags::WithDefaultTimestamp);
        }

        QueryParams {
            consistency: self.consistency,
            flags: flags,
            values: self.values,
            page_size: self.page_size,
            paging_state: self.paging_state,
            serial_consistency: self.serial_consistency,
            timestamp: self.timestamp,
        }

    }
}

pub type QueryBatch = BodyReqBatch;

#[derive(Debug)]
pub struct BatchQueryBuilder {
    batch_type: BatchType,
    queries: Vec<BatchQuery>,
    consistency: Consistency,
    serial_consistency: Option<Consistency>,
    timestamp: Option<i32>,
}

impl BatchQueryBuilder {
    pub fn new() -> BatchQueryBuilder {
        BatchQueryBuilder {
            batch_type: BatchType::Logged,
            queries: vec![],
            consistency: Consistency::One,
            serial_consistency: None,
            timestamp: None,
        }
    }

    pub fn batch_type(mut self, batch_type: BatchType) -> Self {
        self.batch_type = batch_type;
        self
    }

    /// Add a query (non-prepared one)
    pub fn add_query<T: Into<String>>(mut self, query: T, values: Vec<BatchValue>) -> Self {
        self.queries.push(BatchQuery {
            is_prepared: false,
            subject: BatchQuerySubj::QueryString(CStringLong::new(query.into())),
            values: values,
        });
        self
    }

    /// Add a query (prepared one)
    pub fn add_query_prepared(mut self, query_id: CBytesShort, values: Vec<BatchValue>) -> Self {
        self.queries.push(BatchQuery {
            is_prepared: true,
            subject: BatchQuerySubj::PreparedId(query_id),
            values: values,
        });
        self
    }

    pub fn clear_queries(mut self) -> Self {
        self.queries = vec![];
        self
    }

    pub fn consistency(mut self, consistency: Consistency) -> Self {
        self.consistency = consistency;
        self
    }

    pub fn serial_consistency(mut self, serial_consistency: Option<Consistency>) -> Self {
        self.serial_consistency = serial_consistency;
        self
    }

    pub fn timestamp(mut self, timestamp: Option<i32>) -> Self {
        self.timestamp = timestamp;
        self
    }

    pub fn finalize(self) -> CResult<BodyReqBatch> {
        let mut flags = vec![];

        if self.serial_consistency.is_some() {
            flags.push(QueryFlags::WithSerialConsistency);
        }

        if self.timestamp.is_some() {
            flags.push(QueryFlags::WithDefaultTimestamp);
        }

        let with_names_for_values = self.queries
            .iter()
            .all(|q| q.values.iter().all(|v| v.0.is_some()));

        if !with_names_for_values {
            let some_names_for_values = self.queries
                .iter()
                .any(|q| q.values.iter().any(|v| v.0.is_some()));

            if some_names_for_values {
                return Err(CError::General(String::from("Inconsistent query values - mixed \
                                                         with and without names values")));
            }
        }

        if with_names_for_values {
            flags.push(QueryFlags::WithNamesForValues);
        }

        Ok(BodyReqBatch {
            batch_type: self.batch_type,
            queries: self.queries,
            query_flags: flags,
            consistency: self.consistency,
            serial_consistency: self.serial_consistency,
            timestamp: self.timestamp,
        })
    }
}

pub type BatchValue = (Option<CString>, Value);

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

    #[test]
    fn new() {
        let _ = QueryBuilder::new("USE keyspace").finalize();
    }

    #[test]
    fn with_parameters() {
        let _ = QueryBuilder::new("USE keyspace")
            .consistency(Consistency::Two)
            .values(vec![Value::new_null()])
            .with_names(true)
            .page_size(11)
            .paging_state(CBytes::new(vec![1, 2, 3, 4, 5]))
            .serial_consistency(Consistency::One)
            .timestamp(1245678)
            .finalize();
    }
}

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

    #[test]
    fn new() {
        let _ = QueryParamsBuilder::new(Consistency::Two).finalize();
    }

    #[test]
    fn with_parameters() {
        let _ = QueryParamsBuilder::new(Consistency::Two)
            .values(vec![Value::new_null()])
            .with_names(true)
            .page_size(11)
            .paging_state(CBytes::new(vec![1, 2, 3, 4, 5]))
            .serial_consistency(Consistency::One)
            .timestamp(1245678)
            .finalize();
    }
}

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

    #[test]
    fn new() {
        assert!(BatchQueryBuilder::new().finalize().is_ok());
    }

    #[test]
    fn with_parameters() {
        let q = BatchQueryBuilder::new()
            .batch_type(BatchType::Logged)
            .add_query("some query".to_string(), vec![])
            .add_query_prepared(CBytesShort::new(vec![1, 2, 3]), vec![])
            .consistency(Consistency::One)
            .serial_consistency(Some(Consistency::Two))
            .timestamp(Some(1234))
            .finalize();
        assert!(q.is_ok())
    }

    #[test]
    fn clear_queries() {
        let q = BatchQueryBuilder::new()
            .add_query("some query".to_string(), vec![])
            .clear_queries()
            .finalize();
        assert!(q.is_ok())
    }
}
<|endoftext|>


<commit_before>use std::net::{SocketAddr, IpAddr};
use std::io;
use std::time::Duration;

use rand;
use mioco::udp::UdpSocket;
use mioco::tcp::{TcpStream};
use trust_dns::op::{Message, MessageType, ResponseCode, Query, OpCode, Edns};

use utils::{Result, Error, CloneExt, MessageExt, WithTimeout, Future, AsDisplay};
use utils::with_timeout::TcpStreamExt;
use transport::{DnsTransport, DnsMsgTransport};
use validator::{ResponseValidator, DummyValidator};

pub const EDNS_VER: u8 = 0;
pub const EDNS_MAX_PAYLOAD: u16 = 1200;

// Idea: Use DNSSEC to check whether a domain is poisoned by GFW

#[derive(Debug)]
pub enum ErrorKind {
    InvalidId,
    TruncatedBogus(Message),
}
#[derive(Eq, PartialEq, Debug)]
pub enum EdnsMode {
    Disabled,
    Enabled,
    // Required,
}

fn query_core<TTransport, TValidator>(q: Query, mut transport: TTransport, edns_mode: EdnsMode, validator: &mut TValidator) -> Result<Message> where TTransport: DnsMsgTransport, TValidator: ResponseValidator {
    let mut msg : Message = Message::new();
    msg.message_type(MessageType::Query);
    msg.id(rand::random());
    msg.op_code(OpCode::Query);
    msg.recursion_desired(true);
    if edns_mode != EdnsMode::Disabled {
        let mut edns = Edns::new();
        edns.set_version(EDNS_VER);
        edns.set_dnssec_ok(true);
        edns.set_max_payload(EDNS_MAX_PAYLOAD);
        msg.set_edns(edns);
        msg.checking_disabled(true);
    }
    msg.add_query(q);
    debug!("[{}] {}", transport, msg.as_disp());
    try!(transport.send_msg(&msg));
    loop {
        let resp = try!(transport.recv_msg());
        if ! resp.is_resp_for(&msg) {
            warn!("[{}][{}] Invalid response for {}: {:?}",
                  transport, msg.get_id(), msg.as_disp(), resp);
            continue;
        }
        if !validator.is_valid(&resp) {
            if resp.is_truncated() {
                return Err(ErrorKind::TruncatedBogus(msg).into());
            }
            warn!("[{}][{}] Rejected by validator for {}: {}",
                  transport, msg.get_id(), msg.as_disp(), resp.as_disp());
            continue;
        }
        if resp.get_response_code() == ResponseCode::FormErr &&
            edns_mode == EdnsMode::Enabled
        {
            // Maybe the server doesn't implement EDNS?
            return query_core(
                msg.get_queries()[0].clone(), transport, EdnsMode::Disabled, validator,
            );
        }
        debug!("[{}] {} -> {}", transport, resp.get_queries()[0].as_disp(), resp.as_disp());
        return Ok(resp);
    }
}
fn get_bind_addr(target: &IpAddr) -> SocketAddr {
    SocketAddr::new(match *target {
        IpAddr::V4(_) => "0.0.0.0",
        IpAddr::V6(_) => "::0",
    }.parse().unwrap(), 0)
}
pub fn query(q: Query, addr: IpAddr, timeout: Duration) -> Result<Message> {
    query_with_validator(q, addr, timeout, &mut DummyValidator)
}
pub fn query_with_validator<T: ResponseValidator>(q: Query, addr: IpAddr, timeout: Duration, validator: &mut T) -> Result<Message> {
    let target = SocketAddr::new(addr, 53);
    let mut transport = try!(UdpSocket::bound(&get_bind_addr(&addr))).with_timeout(timeout);
    macro_rules! query_tcp {
        ($q: expr) => {{
            TcpStream::connect_with_timeout(&target, timeout)
            .map_err(|e| e.into())
            .and_then(|stream| stream.set_nodelay(true)
                                     .or_else(|_| Ok(()))
                                     .map(|_| stream))
            .and_then(|stream| query_core(
                $q.clone(),
                stream.with_timeout(timeout).bound(Some(&target)),
                EdnsMode::Enabled,
                validator,
            ))
        }}
    };
    match query_core(q, transport.bound(Some(&target)), EdnsMode::Enabled, validator) {
        Ok(msg) => {
            if msg.is_truncated() {
                // Try TCP
                Ok(query_tcp!(&msg.get_queries()[0]).unwrap_or(msg))
            } else {
                Ok(msg)
            }
        },
        Err(Error::Query(ErrorKind::TruncatedBogus(msg))) => {
            query_tcp!(&msg.get_queries()[0])
        },
        Err(e) => Err(e),
    }
}
pub fn query_multiple_handle_futures(futures: &mut Vec<Future<Result<Message>>>) -> Result<Message> {
    if futures.is_empty() {
        return Err(io::ErrorKind::InvalidInput.into());
    }
    loop {
        let result_index = Future::wait_any(futures);
        let result = futures.remove(result_index).consume();
        let mut should_return = futures.is_empty();
        match result {
            Ok(ref msg) => {
                match msg.get_response_code() {
                    ResponseCode::ServFail |
                    ResponseCode::NotImp |
                    ResponseCode::Refused => {
                        // Fall below to wait for remaining servers, if any
                    },
                    _ => { should_return = true; },
                }
            },
            Err(Error::Io(ref err)) if err.kind() == io::ErrorKind::TimedOut => {
                debug!("Timed out");
            },
            Err(ref err) => {
                debug!("Error {:?}", err);
            },
        };
        if should_return {
            return result;
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use trust_dns::op::*;
    use trust_dns::rr::*;

    use super::*;
    use ::mioco_config_start;

    #[test]
    fn simple_query() {
        mioco_config_start(|| {
            let mut q = Query::new();
            q.name(Name::parse("www.google.com", Some(&Name::root())).unwrap());
            let result = query(q, "8.8.8.8".parse().unwrap(), Duration::from_secs(5)).unwrap();
            assert!(result.get_answers().len() > 0);
        }).unwrap();
    }
}
<commit_msg>Separate validation error from timeout<commit_after>use std::net::{SocketAddr, IpAddr};
use std::io;
use std::time::Duration;

use rand;
use mioco::udp::UdpSocket;
use mioco::tcp::{TcpStream};
use trust_dns::op::{Message, MessageType, ResponseCode, Query, OpCode, Edns};

use utils::{Result, Error, CloneExt, MessageExt, WithTimeout, Future, AsDisplay};
use utils::with_timeout::TcpStreamExt;
use transport::{DnsTransport, DnsMsgTransport};
use validator::{ResponseValidator, DummyValidator};

pub const EDNS_VER: u8 = 0;
pub const EDNS_MAX_PAYLOAD: u16 = 1200;

// Idea: Use DNSSEC to check whether a domain is poisoned by GFW

#[derive(Debug)]
pub enum ErrorKind {
    InvalidId,
    TruncatedBogus(Message),
    ValidationFailure(Message),
}
#[derive(Eq, PartialEq, Debug)]
pub enum EdnsMode {
    Disabled,
    Enabled,
    // Required,
}

fn query_core<TTransport, TValidator>(q: Query, mut transport: TTransport, edns_mode: EdnsMode, validator: &mut TValidator) -> Result<Message> where TTransport: DnsMsgTransport, TValidator: ResponseValidator {
    let mut msg : Message = Message::new();
    msg.message_type(MessageType::Query);
    msg.id(rand::random());
    msg.op_code(OpCode::Query);
    msg.recursion_desired(true);
    if edns_mode != EdnsMode::Disabled {
        let mut edns = Edns::new();
        edns.set_version(EDNS_VER);
        edns.set_dnssec_ok(true);
        edns.set_max_payload(EDNS_MAX_PAYLOAD);
        msg.set_edns(edns);
        msg.checking_disabled(true);
    }
    msg.add_query(q);
    debug!("[{}] {}", transport, msg.as_disp());
    try!(transport.send_msg(&msg));
    let mut last_bogus_msg: Option<Message> = None;
    loop {
        let resp = try!(match transport.recv_msg() {
            Ok(x) => Ok(x),
            Err(e) => Err(match e {
                Error::Io(e) => {
                    match e.kind() {
                        io::ErrorKind::TimedOut if last_bogus_msg.is_some() => {
                            ErrorKind::ValidationFailure(
                                last_bogus_msg.take().unwrap()
                            ).into()
                        },
                        _ => Error::Io(e),
                    }
                },
                _ => e,
            }),
        });
        if ! resp.is_resp_for(&msg) {
            warn!("[{}][{}] Invalid response for {}: {:?}",
                  transport, msg.get_id(), msg.as_disp(), resp);
            continue;
        }
        if !validator.is_valid(&resp) {
            if resp.is_truncated() {
                return Err(ErrorKind::TruncatedBogus(msg).into());
            }
            warn!("[{}][{}] Rejected by validator for {}: {}",
                  transport, msg.get_id(), msg.as_disp(), resp.as_disp());
            last_bogus_msg = Some(resp);
            continue;
        }
        if resp.get_response_code() == ResponseCode::FormErr &&
            edns_mode == EdnsMode::Enabled
        {
            // Maybe the server doesn't implement EDNS?
            return query_core(
                msg.get_queries()[0].clone(), transport, EdnsMode::Disabled, validator,
            );
        }
        debug!("[{}] {} -> {}", transport, resp.get_queries()[0].as_disp(), resp.as_disp());
        return Ok(resp);
    }
}
fn get_bind_addr(target: &IpAddr) -> SocketAddr {
    SocketAddr::new(match *target {
        IpAddr::V4(_) => "0.0.0.0",
        IpAddr::V6(_) => "::0",
    }.parse().unwrap(), 0)
}
pub fn query(q: Query, addr: IpAddr, timeout: Duration) -> Result<Message> {
    query_with_validator(q, addr, timeout, &mut DummyValidator)
}
pub fn query_with_validator<T: ResponseValidator>(q: Query, addr: IpAddr, timeout: Duration, validator: &mut T) -> Result<Message> {
    let target = SocketAddr::new(addr, 53);
    let mut transport = try!(UdpSocket::bound(&get_bind_addr(&addr))).with_timeout(timeout);
    macro_rules! query_tcp {
        ($q: expr) => {{
            TcpStream::connect_with_timeout(&target, timeout)
            .map_err(|e| e.into())
            .and_then(|stream| stream.set_nodelay(true)
                                     .or_else(|_| Ok(()))
                                     .map(|_| stream))
            .and_then(|stream| query_core(
                $q.clone(),
                stream.with_timeout(timeout).bound(Some(&target)),
                EdnsMode::Enabled,
                validator,
            ))
        }}
    };
    match query_core(q, transport.bound(Some(&target)), EdnsMode::Enabled, validator) {
        Ok(msg) => {
            if msg.is_truncated() {
                // Try TCP
                Ok(query_tcp!(&msg.get_queries()[0]).unwrap_or(msg))
            } else {
                Ok(msg)
            }
        },
        Err(Error::Query(ErrorKind::TruncatedBogus(msg))) => {
            query_tcp!(&msg.get_queries()[0])
        },
        Err(e) => Err(e),
    }
}
pub fn query_multiple_handle_futures(futures: &mut Vec<Future<Result<Message>>>) -> Result<Message> {
    if futures.is_empty() {
        return Err(io::ErrorKind::InvalidInput.into());
    }
    loop {
        let result_index = Future::wait_any(futures);
        let result = futures.remove(result_index).consume();
        let mut should_return = futures.is_empty();
        match result {
            Ok(ref msg) => {
                match msg.get_response_code() {
                    ResponseCode::ServFail |
                    ResponseCode::NotImp |
                    ResponseCode::Refused => {
                        // Fall below to wait for remaining servers, if any
                    },
                    _ => { should_return = true; },
                }
            },
            Err(Error::Io(ref err)) if err.kind() == io::ErrorKind::TimedOut => {
                debug!("Timed out");
            },
            Err(ref err) => {
                debug!("Error {:?}", err);
            },
        };
        if should_return {
            return result;
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use trust_dns::op::*;
    use trust_dns::rr::*;

    use super::*;
    use ::mioco_config_start;

    #[test]
    fn simple_query() {
        mioco_config_start(|| {
            let mut q = Query::new();
            q.name(Name::parse("www.google.com", Some(&Name::root())).unwrap());
            let result = query(q, "8.8.8.8".parse().unwrap(), Duration::from_secs(5)).unwrap();
            assert!(result.get_answers().len() > 0);
        }).unwrap();
    }
}
<|endoftext|>


<commit_before>//! A control flow graph represented as mappings of extended basic blocks to their predecessors
//! and successors. Successors are represented as extended basic blocks while predecessors are
//! represented by basic blocks.
//! BasicBlocks are denoted by tuples of EBB and branch/jump instructions. Each predecessor
//! tuple corresponds to the end of a basic block.
//!
//! ```c
//!     Ebb0:
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         brz vx, Ebb1 ; end of basic block
//!
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         jmp Ebb2     ; end of basic block
//! ```
//!
//! Here Ebb1 and Ebb2 would each have a single predecessor denoted as (Ebb0, `brz vx, Ebb1`)
//! and (Ebb0, `jmp Ebb2`) respectively.

use ir::Function;
use ir::entities::{Inst, Ebb};
use ir::instructions::BranchInfo;
use entity_map::EntityMap;

/// A basic block denoted by its enclosing Ebb and last instruction.
pub type BasicBlock = (Ebb, Inst);

/// A container for the successors and predecessors of some Ebb.
#[derive(Debug)]
pub struct CFGNode {
    pub successors: Vec<Ebb>,
    pub predecessors: Vec<BasicBlock>,
}

impl CFGNode {
    pub fn new() -> CFGNode {
        CFGNode {
            successors: Vec::new(),
            predecessors: Vec::new(),
        }
    }
}

/// The Control Flow Graph maintains a mapping of ebbs to their predecessors
/// and successors where predecessors are basic blocks and successors are
/// extended basic blocks.
#[derive(Debug)]
pub struct ControlFlowGraph {
    data: EntityMap<Ebb, CFGNode>,
}

impl ControlFlowGraph {
    /// During initialization mappings will be generated for any existing
    /// blocks within the CFG's associated function.
    pub fn new(func: &Function) -> ControlFlowGraph {
        let mut cfg = ControlFlowGraph { data: EntityMap::new() };

        // Even ebbs without predecessors should show up in the CFG, albeit
        // with no entires.
        for _ in &func.layout {
            cfg.push_ebb();
        }

        for ebb in &func.layout {
            for inst in func.layout.ebb_insts(ebb) {
                match func.dfg[inst].analyze_branch() {
                    BranchInfo::SingleDest(dest, _) => {
                        cfg.add_successor(ebb, dest);
                        cfg.add_predecessor(dest, (ebb, inst));
                    }
                    BranchInfo::Table(jt) => {
                        for (_, dest) in func.jump_tables[jt].entries() {
                            cfg.add_successor(ebb, dest);
                            cfg.add_predecessor(dest, (ebb, inst));
                        }
                    }
                    BranchInfo::NotABranch => {}
                }
            }
        }
        cfg
    }

    pub fn push_ebb(&mut self) {
        self.data.push(CFGNode::new());
    }

    pub fn add_successor(&mut self, from: Ebb, to: Ebb) {
        self.data[from].successors.push(to);
    }

    pub fn add_predecessor(&mut self, ebb: Ebb, predecessor: BasicBlock) {
        self.data[ebb].predecessors.push(predecessor);
    }

    pub fn get_predecessors(&self, ebb: Ebb) -> &Vec<BasicBlock> {
        &self.data[ebb].predecessors
    }

    pub fn get_successors(&self, ebb: Ebb) -> &Vec<Ebb> {
        &self.data[ebb].successors
    }

    pub fn postorder_ebbs(&self, entry: Ebb) -> Vec<Ebb> {
        let mut stack_a = vec![entry];
        let mut stack_b = Vec::new();
        while stack_a.len() > 0 {
            let cur = stack_a.pop().unwrap();
            for child in &self.data[cur].successors {
                if *child != cur && !stack_a.contains(child) {
                    stack_a.push(child.clone());
                }
            }
            stack_b.push(cur);
        }
        stack_b
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn predecessors_iter(&self) -> CFGPredecessorsIter {
        CFGPredecessorsIter {
            cur: 0,
            cfg: &self,
        }
    }
}

/// Iterate through every mapping of ebb to predecessors in the CFG
pub struct CFGPredecessorsIter<'a> {
    cfg: &'a ControlFlowGraph,
    cur: usize,
}

impl<'a> Iterator for CFGPredecessorsIter<'a> {
    type Item = (Ebb, &'a Vec<BasicBlock>);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur < self.cfg.len() {
            let ebb = Ebb::with_number(self.cur as u32).unwrap();
            let bbs = self.cfg.get_predecessors(ebb);
            self.cur += 1;
            Some((ebb, bbs))
        } else {
            None
        }
    }
}

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

    use test_utils::make_inst;

    #[test]
    fn empty() {
        let func = Function::new();
        let cfg = ControlFlowGraph::new(&func);
        assert_eq!(None, cfg.predecessors_iter().next());
    }

    #[test]
    fn no_predecessors() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let cfg = ControlFlowGraph::new(&func);
        let nodes = cfg.predecessors_iter().collect::<Vec<_>>();
        assert_eq!(nodes.len(), 3);

        let mut fun_ebbs = func.layout.ebbs();
        for (ebb, predecessors) in nodes {
            assert_eq!(ebb, fun_ebbs.next().unwrap());
            assert_eq!(predecessors.len(), 0);
            assert_eq!(predecessors.len(), 0);
            assert_eq!(cfg.get_successors(ebb).len(), 0);
        }
    }

    #[test]
    fn branches_and_jumps() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let br_ebb0_ebb2 = make_inst::branch(&mut func, ebb2);
        func.layout.append_inst(br_ebb0_ebb2, ebb0);

        let jmp_ebb0_ebb1 = make_inst::jump(&mut func, ebb1);
        func.layout.append_inst(jmp_ebb0_ebb1, ebb0);

        let br_ebb1_ebb1 = make_inst::branch(&mut func, ebb1);
        func.layout.append_inst(br_ebb1_ebb1, ebb1);

        let jmp_ebb1_ebb2 = make_inst::jump(&mut func, ebb2);
        func.layout.append_inst(jmp_ebb1_ebb2, ebb1);

        let cfg = ControlFlowGraph::new(&func);

        let ebb0_predecessors = cfg.get_predecessors(ebb0);
        let ebb1_predecessors = cfg.get_predecessors(ebb1);
        let ebb2_predecessors = cfg.get_predecessors(ebb2);

        let ebb0_successors = cfg.get_successors(ebb0);
        let ebb1_successors = cfg.get_successors(ebb1);
        let ebb2_successors = cfg.get_successors(ebb2);

        assert_eq!(ebb0_predecessors.len(), 0);
        assert_eq!(ebb1_predecessors.len(), 2);
        assert_eq!(ebb2_predecessors.len(), 2);

        assert_eq!(ebb1_predecessors.contains(&(ebb0, jmp_ebb0_ebb1)), true);
        assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
        assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), true);
        assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

        assert_eq!(ebb0_successors.len(), 2);
        assert_eq!(ebb1_successors.len(), 2);
        assert_eq!(ebb2_successors.len(), 0);

        assert_eq!(ebb0_successors.contains(&ebb1), true);
        assert_eq!(ebb0_successors.contains(&ebb2), true);
        assert_eq!(ebb1_successors.contains(&ebb1), true);
        assert_eq!(ebb1_successors.contains(&ebb2), true);
    }

    #[test]
    fn postorder_traversal() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        let ebb3 = func.dfg.make_ebb();
        let ebb4 = func.dfg.make_ebb();
        let ebb5 = func.dfg.make_ebb();

        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);
        func.layout.append_ebb(ebb3);
        func.layout.append_ebb(ebb4);
        func.layout.append_ebb(ebb5);

        let br_ebb0_ebb1 = make_inst::branch(&mut func, ebb1);
        func.layout.append_inst(br_ebb0_ebb1, ebb0);

        let jmp_ebb0_ebb2 = make_inst::jump(&mut func, ebb2);
        func.layout.append_inst(jmp_ebb0_ebb2, ebb0);

        let br_ebb2_ebb2 = make_inst::branch(&mut func, ebb2);
        func.layout.append_inst(br_ebb2_ebb2, ebb2);

        let br_ebb2_ebb1 = make_inst::branch(&mut func, ebb1);
        func.layout.append_inst(br_ebb2_ebb1, ebb2);

        let jmp_ebb1_ebb3 = make_inst::jump(&mut func, ebb3);
        func.layout.append_inst(jmp_ebb1_ebb3, ebb1);

        let br_ebb2_ebb4 = make_inst::branch(&mut func, ebb4);
        func.layout.append_inst(br_ebb2_ebb4, ebb2);

        let jmp_ebb2_ebb5 = make_inst::jump(&mut func, ebb5);
        func.layout.append_inst(jmp_ebb2_ebb5, ebb2);

        let cfg = ControlFlowGraph::new(&func);
        assert_eq!(cfg.postorder_ebbs(func.layout.entry_block().unwrap()),
                   vec![ebb0, ebb2, ebb5, ebb4, ebb1, ebb3]);
    }

    #[test]
    fn loop_edge() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        let ebb3 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);
        func.layout.append_ebb(ebb3);

        let jmp_ebb0_ebb1 = make_inst::jump(&mut func, ebb1);
        let br_ebb1_ebb3 = make_inst::branch(&mut func, ebb3);
        let jmp_ebb1_ebb2 = make_inst::jump(&mut func, ebb2);
        let jmp_ebb2_ebb3 = make_inst::jump(&mut func, ebb3);

        func.layout.append_inst(jmp_ebb0_ebb1, ebb0);
        func.layout.append_inst(br_ebb1_ebb3, ebb1);
        func.layout.append_inst(jmp_ebb1_ebb2, ebb1);
        func.layout.append_inst(jmp_ebb2_ebb3, ebb2);

        let cfg = ControlFlowGraph::new(&func);
        assert_eq!(cfg.postorder_ebbs(func.layout.entry_block().unwrap()),
                   vec![ebb0, ebb1, ebb2, ebb3]);
    }
}
<commit_msg>make postorder_ebbs into actually reverse_postorder_ebbs.<commit_after>//! A control flow graph represented as mappings of extended basic blocks to their predecessors
//! and successors. Successors are represented as extended basic blocks while predecessors are
//! represented by basic blocks.
//! BasicBlocks are denoted by tuples of EBB and branch/jump instructions. Each predecessor
//! tuple corresponds to the end of a basic block.
//!
//! ```c
//!     Ebb0:
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         brz vx, Ebb1 ; end of basic block
//!
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         jmp Ebb2     ; end of basic block
//! ```
//!
//! Here Ebb1 and Ebb2 would each have a single predecessor denoted as (Ebb0, `brz vx, Ebb1`)
//! and (Ebb0, `jmp Ebb2`) respectively.

use ir::Function;
use ir::entities::{Inst, Ebb};
use ir::instructions::InstructionData;
use entity_map::EntityMap;
use std::collections::{HashSet, BTreeMap};

/// A basic block denoted by its enclosing Ebb and last instruction.
pub type BasicBlock = (Ebb, Inst);

/// A container for the successors and predecessors of some Ebb.
#[derive(Debug)]
pub struct CFGNode {
    pub successors: Vec<Ebb>,
    pub predecessors: Vec<BasicBlock>,
}

impl CFGNode {
    pub fn new() -> CFGNode {
        CFGNode {
            successors: Vec::new(),
            predecessors: Vec::new(),
        }
    }
}

/// The Control Flow Graph maintains a mapping of ebbs to their predecessors
/// and successors where predecessors are basic blocks and successors are
/// extended basic blocks.
#[derive(Debug)]
pub struct ControlFlowGraph {
    entry_block: Option<Ebb>,
    data: EntityMap<Ebb, CFGNode>,
}

impl ControlFlowGraph {
    /// During initialization mappings will be generated for any existing
    /// blocks within the CFG's associated function.
    pub fn new(func: &Function) -> ControlFlowGraph {

        let mut cfg = ControlFlowGraph {
            data: EntityMap::new(),
            entry_block: func.layout.entry_block(),
        };

        // Even ebbs without predecessors should show up in the CFG, albeit
        // with no entires.
        for _ in &func.layout {
            cfg.push_ebb();
        }

        for ebb in &func.layout {
            for inst in func.layout.ebb_insts(ebb) {
                match func.dfg[inst] {
                    InstructionData::Branch { ty: _, opcode: _, ref data } => {
                        cfg.add_successor(ebb, data.destination);
                        cfg.add_predecessor(data.destination, (ebb, inst));
                    }
                    InstructionData::Jump { ty: _, opcode: _, ref data } => {
                        cfg.add_successor(ebb, data.destination);
                        cfg.add_predecessor(data.destination, (ebb, inst));
                    }
                    _ => (),
                }
            }
        }
        cfg
    }

    pub fn push_ebb(&mut self) {
        self.data.push(CFGNode::new());
    }

    pub fn add_successor(&mut self, from: Ebb, to: Ebb) {
        self.data[from].successors.push(to);
    }

    pub fn add_predecessor(&mut self, ebb: Ebb, predecessor: BasicBlock) {
        self.data[ebb].predecessors.push(predecessor);
    }

    pub fn get_predecessors(&self, ebb: Ebb) -> &Vec<BasicBlock> {
        &self.data[ebb].predecessors
    }

    pub fn get_successors(&self, ebb: Ebb) -> &Vec<Ebb> {
        &self.data[ebb].successors
    }

    /// Return ebbs in reverse postorder along with a mapping of
    /// the ebb to its order of visitation.
    pub fn reverse_postorder_ebbs(&self) -> BTreeMap<Ebb, usize> {
        let entry_block = match self.entry_block {
            None => {
                return BTreeMap::new();
            }
            Some(eb) => eb,
        };
        let mut seen = HashSet::new();
        let mut stack_a = vec![entry_block];
        let mut finished = BTreeMap::new();
        while stack_a.len() > 0 {
            let cur = stack_a.pop().unwrap();
            for child in &self.data[cur].successors {
                if *child != cur && !seen.contains(&child) {
                    seen.insert(child);
                    stack_a.push(child.clone());
                }
            }
            let index = finished.len();
            finished.insert(cur, index);
        }
        finished
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn predecessors_iter(&self) -> CFGPredecessorsIter {
        CFGPredecessorsIter {
            cur: 0,
            cfg: &self,
        }
    }
}

/// Iterate through every mapping of ebb to predecessors in the CFG
pub struct CFGPredecessorsIter<'a> {
    cfg: &'a ControlFlowGraph,
    cur: usize,
}

impl<'a> Iterator for CFGPredecessorsIter<'a> {
    type Item = (Ebb, &'a Vec<BasicBlock>);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur < self.cfg.len() {
            let ebb = Ebb::with_number(self.cur as u32).unwrap();
            let bbs = self.cfg.get_predecessors(ebb);
            self.cur += 1;
            Some((ebb, bbs))
        } else {
            None
        }
    }
}

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

    use test_utils::make_inst;

    #[test]
    fn empty() {
        let func = Function::new();
        let cfg = ControlFlowGraph::new(&func);
        assert_eq!(None, cfg.predecessors_iter().next());
    }

    #[test]
    fn no_predecessors() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let cfg = ControlFlowGraph::new(&func);
        let nodes = cfg.predecessors_iter().collect::<Vec<_>>();
        assert_eq!(nodes.len(), 3);

        let mut fun_ebbs = func.layout.ebbs();
        for (ebb, predecessors) in nodes {
            assert_eq!(ebb, fun_ebbs.next().unwrap());
            assert_eq!(predecessors.len(), 0);
            assert_eq!(predecessors.len(), 0);
            assert_eq!(cfg.get_successors(ebb).len(), 0);
        }
    }

    #[test]
    fn branches_and_jumps() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let br_ebb0_ebb2 = make_inst::branch(&mut func, ebb2);
        func.layout.append_inst(br_ebb0_ebb2, ebb0);

        let jmp_ebb0_ebb1 = make_inst::jump(&mut func, ebb1);
        func.layout.append_inst(jmp_ebb0_ebb1, ebb0);

        let br_ebb1_ebb1 = make_inst::branch(&mut func, ebb1);
        func.layout.append_inst(br_ebb1_ebb1, ebb1);

        let jmp_ebb1_ebb2 = make_inst::jump(&mut func, ebb2);
        func.layout.append_inst(jmp_ebb1_ebb2, ebb1);

        let cfg = ControlFlowGraph::new(&func);

        let ebb0_predecessors = cfg.get_predecessors(ebb0);
        let ebb1_predecessors = cfg.get_predecessors(ebb1);
        let ebb2_predecessors = cfg.get_predecessors(ebb2);

        let ebb0_successors = cfg.get_successors(ebb0);
        let ebb1_successors = cfg.get_successors(ebb1);
        let ebb2_successors = cfg.get_successors(ebb2);

        assert_eq!(ebb0_predecessors.len(), 0);
        assert_eq!(ebb1_predecessors.len(), 2);
        assert_eq!(ebb2_predecessors.len(), 2);

        assert_eq!(ebb1_predecessors.contains(&(ebb0, jmp_ebb0_ebb1)), true);
        assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
        assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), true);
        assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

        assert_eq!(ebb0_successors.len(), 2);
        assert_eq!(ebb1_successors.len(), 2);
        assert_eq!(ebb2_successors.len(), 0);

        assert_eq!(ebb0_successors.contains(&ebb1), true);
        assert_eq!(ebb0_successors.contains(&ebb2), true);
        assert_eq!(ebb1_successors.contains(&ebb1), true);
        assert_eq!(ebb1_successors.contains(&ebb2), true);
    }

    #[test]
    fn postorder_traversal() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        let ebb3 = func.dfg.make_ebb();
        let ebb4 = func.dfg.make_ebb();
        let ebb5 = func.dfg.make_ebb();

        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);
        func.layout.append_ebb(ebb3);
        func.layout.append_ebb(ebb4);
        func.layout.append_ebb(ebb5);

        let br_ebb0_ebb1 = make_inst::branch(&mut func, ebb1);
        func.layout.append_inst(br_ebb0_ebb1, ebb0);

        let jmp_ebb0_ebb2 = make_inst::jump(&mut func, ebb2);
        func.layout.append_inst(jmp_ebb0_ebb2, ebb0);

        let br_ebb2_ebb2 = make_inst::branch(&mut func, ebb2);
        func.layout.append_inst(br_ebb2_ebb2, ebb2);

        let br_ebb2_ebb1 = make_inst::branch(&mut func, ebb1);
        func.layout.append_inst(br_ebb2_ebb1, ebb2);

        let jmp_ebb1_ebb3 = make_inst::jump(&mut func, ebb3);
        func.layout.append_inst(jmp_ebb1_ebb3, ebb1);

        let br_ebb2_ebb4 = make_inst::branch(&mut func, ebb4);
        func.layout.append_inst(br_ebb2_ebb4, ebb2);

        let jmp_ebb2_ebb5 = make_inst::jump(&mut func, ebb5);
        func.layout.append_inst(jmp_ebb2_ebb5, ebb2);

        let cfg = ControlFlowGraph::new(&func);
        let mut postorder = vec![ebb3, ebb1, ebb4, ebb5, ebb2, ebb0];
        postorder.reverse();
        for (ebb, key) in cfg.reverse_postorder_ebbs() {
            assert_eq!(ebb, postorder[key]);
        }
    }

    #[test]
    fn loops_one() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        let ebb3 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);
        func.layout.append_ebb(ebb3);

        let jmp_ebb0_ebb1 = make_inst::jump(&mut func, ebb1);
        let br_ebb1_ebb3 = make_inst::branch(&mut func, ebb3);
        let jmp_ebb1_ebb2 = make_inst::jump(&mut func, ebb2);
        let jmp_ebb2_ebb3 = make_inst::jump(&mut func, ebb3);

        func.layout.append_inst(jmp_ebb0_ebb1, ebb0);
        func.layout.append_inst(br_ebb1_ebb3, ebb1);
        func.layout.append_inst(jmp_ebb1_ebb2, ebb1);
        func.layout.append_inst(jmp_ebb2_ebb3, ebb2);

        let cfg = ControlFlowGraph::new(&func);
        let mut postorder = vec![ebb3, ebb2, ebb1, ebb0];
        postorder.reverse();
        for (ebb, key) in cfg.reverse_postorder_ebbs() {
            assert_eq!(ebb, postorder[key]);
        }
    }

    #[test]
    fn loops_two() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        let ebb3 = func.dfg.make_ebb();
        let ebb4 = func.dfg.make_ebb();
        let ebb5 = func.dfg.make_ebb();

        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);
        func.layout.append_ebb(ebb3);
        func.layout.append_ebb(ebb4);
        func.layout.append_ebb(ebb5);

        let jmp_ebb0_ebb1 = make_inst::jump(&mut func, ebb1);
        let jmp_ebb0_ebb2 = make_inst::jump(&mut func, ebb2);
        let jmp_ebb1_ebb3 = make_inst::jump(&mut func, ebb3);
        let br_ebb2_ebb4 = make_inst::jump(&mut func, ebb4);
        let jmp_ebb2_ebb5 = make_inst::jump(&mut func, ebb5);
        let jmp_ebb3_ebb4 = make_inst::jump(&mut func, ebb4);
        let br_ebb4_ebb3 = make_inst::branch(&mut func, ebb3);
        let jmp_ebb4_ebb5 = make_inst::jump(&mut func, ebb5);
        let jmp_ebb5_ebb4 = make_inst::jump(&mut func, ebb4);

        func.layout.append_inst(jmp_ebb0_ebb1, ebb0);
        func.layout.append_inst(jmp_ebb0_ebb2, ebb0);
        func.layout.append_inst(jmp_ebb1_ebb3, ebb1);
        func.layout.append_inst(br_ebb2_ebb4, ebb2);
        func.layout.append_inst(jmp_ebb2_ebb5, ebb2);
        func.layout.append_inst(jmp_ebb3_ebb4, ebb3);
        func.layout.append_inst(br_ebb4_ebb3, ebb4);
        func.layout.append_inst(jmp_ebb4_ebb5, ebb4);
        func.layout.append_inst(jmp_ebb5_ebb4, ebb5);

        let cfg = ControlFlowGraph::new(&func);
        let mut postorder = vec![ebb1, ebb3, ebb4, ebb5, ebb2, ebb0];
        postorder.reverse();
        for (ebb, key) in cfg.reverse_postorder_ebbs() {
            assert_eq!(ebb, postorder[key]);
        }
    }
}
<|endoftext|>


<commit_before>#[cfg(feature = "serde-serialize")]
use serde::{Deserialize, Serialize};

use num::One;
use alga::general::ComplexField;

use crate::allocator::Allocator;
use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, SquareMatrix, Vector};
use crate::constraint::{SameNumberOfRows, ShapeConstraint};
use crate::dimension::{Dim, DimAdd, DimSum, DimDiff, DimSub, Dynamic, U1};
use crate::storage::{Storage, StorageMut};

/// The Cholesky decomposition of a symmetric-definite-positive matrix.
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(serialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Serialize"))
)]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(deserialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Deserialize<'de>"))
)]
#[derive(Clone, Debug)]
pub struct Cholesky<N: ComplexField, D: Dim>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    chol: MatrixN<N, D>,
}

impl<N: ComplexField, D: Dim> Copy for Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
    MatrixN<N, D>: Copy,
{
}

impl<N: ComplexField, D: DimSub<Dynamic>> Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of `matrix`.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn new(mut matrix: MatrixN<N, D>) -> Option<Self> {
        assert!(matrix.is_square(), "The input matrix must be square.");

        let n = matrix.nrows();

        for j in 0..n {
            for k in 0..j {
                let factor = unsafe { -*matrix.get_unchecked((j, k)) };

                let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
                let mut col_j = col_j.rows_range_mut(j..);
                let col_k = col_k.rows_range(j..);

                col_j.axpy(factor.conjugate(), &col_k, N::one());
            }

            let diag = unsafe { *matrix.get_unchecked((j, j)) };
            if !diag.is_zero() {
                if let Some(denom) = diag.try_sqrt() {
                    unsafe {
                        *matrix.get_unchecked_mut((j, j)) = denom;
                    }

                    let mut col = matrix.slice_range_mut(j + 1.., j);
                    col /= denom;
                    continue;
                }
            }

            // The diagonal element is either zero or its square root could not
            // be taken (e.g. for negative real numbers).
            return None;
        }

        Some(Cholesky { chol: matrix })
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// upper-triangular part filled with zeros.
    pub fn unpack(mut self) -> MatrixN<N, D> {
        self.chol.fill_upper_triangle(N::zero(), 1);
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// The values of the strict upper-triangular part are garbage and should be ignored by further
    /// computations.
    pub fn unpack_dirty(self) -> MatrixN<N, D> {
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// uppen-triangular part filled with zeros.
    pub fn l(&self) -> MatrixN<N, D> {
        self.chol.lower_triangle()
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// This is an allocation-less version of `self.l()`. The values of the strict upper-triangular
    /// part are garbage and should be ignored by further computations.
    pub fn l_dirty(&self) -> &MatrixN<N, D> {
        &self.chol
    }

    /// Solves the system `self * x = b` where `self` is the decomposed matrix and `x` the unknown.
    ///
    /// The result is stored on `b`.
    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&self, b: &mut Matrix<N, R2, C2, S2>)
    where
        S2: StorageMut<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let _ = self.chol.solve_lower_triangular_mut(b);
        let _ = self.chol.ad_solve_lower_triangular_mut(b);
    }

    /// Returns the solution of the system `self * x = b` where `self` is the decomposed matrix and
    /// `x` the unknown.
    pub fn solve<R2: Dim, C2: Dim, S2>(&self, b: &Matrix<N, R2, C2, S2>) -> MatrixMN<N, R2, C2>
    where
        S2: Storage<N, R2, C2>,
        DefaultAllocator: Allocator<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let mut res = b.clone_owned();
        self.solve_mut(&mut res);
        res
    }

    /// Computes the inverse of the decomposed matrix.
    pub fn inverse(&self) -> MatrixN<N, D> {
        let shape = self.chol.data.shape();
        let mut res = MatrixN::identity_generic(shape.0, shape.1);

        self.solve_mut(&mut res);
        res
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (v * v.adjoint())`.
    #[inline]
    pub fn rank_one_update<R2: Dim, S2>(&mut self, x: &Vector<N, R2, S2>, sigma: N::RealField)
    where
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, R2, U1>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        Self::xx_rank_one_update(&mut self.chol, &mut x.clone_owned(), sigma)
    }

    /// Updates the decomposition such that we get the decomposition of a matrix with the given column `col` in the `j`th position.
    /// Since the matrix is square, an identical row will be added in the `j`th row.
    pub fn insert_column<R2, S2>(
        &self,
        j: usize,
        col: Vector<N, R2, S2>,
    ) -> Cholesky<N, DimSum<D, U1>>
    where
        D: DimAdd<U1>,
        R2: Dim,
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, DimSum<D, U1>, DimSum<D, U1>> + Allocator<N, R2>,
        ShapeConstraint: SameNumberOfRows<R2, DimSum<D, U1>>,
    {
        let mut col = col.into_owned();
        // for an explanation of the formulas, see https://en.wikipedia.org/wiki/Cholesky_decomposition#Updating_the_decomposition
        let n = col.nrows();
        assert_eq!(n, self.chol.nrows() + 1, "The new column must have the size of the factored matrix plus one.");
        assert!(j < n, "j needs to be within the bound of the new matrix.");

        // loads the data into a new matrix with an additional jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.add(U1), self.chol.data.shape().1.add(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j + 1..).copy_from(&self.chol.slice_range(..j, j..));
        chol.slice_range_mut(j + 1.., ..j).copy_from(&self.chol.slice_range(j.., ..j));
        chol.slice_range_mut(j + 1.., j + 1..).copy_from(&self.chol.slice_range(j.., j..));

        // update the jth row
        let top_left_corner = self.chol.slice_range(..j, ..j);

        let col_j = col[j];
        let (mut new_rowj_adjoint, mut new_colj) = col.rows_range_pair_mut(..j, j + 1..);
        assert!(top_left_corner.solve_lower_triangular_mut(&mut new_rowj_adjoint), "Cholesky::insert_column : Unable to solve lower triangular system!");

        new_rowj_adjoint.adjoint_to(&mut chol.slice_range_mut(j, ..j));

        // update the center element
        let center_element = N::sqrt(col_j - N::from_real(new_rowj_adjoint.norm_squared()));
        chol[(j, j)] = center_element;

        // update the jth column
        let bottom_left_corner = self.chol.slice_range(j.., ..j);
        // new_colj = (col_jplus - bottom_left_corner * new_rowj.adjoint()) / center_element;
        new_colj.gemm(-N::one() / center_element, &bottom_left_corner, &new_rowj_adjoint, N::one() / center_element);
        chol.slice_range_mut(j + 1.., j).copy_from(&new_colj);

        // update the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j + 1.., j + 1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut new_colj, -N::RealField::one());

        Cholesky { chol }
    }

    /// Updates the decomposition such that we get the decomposition of the factored matrix with its `j`th column removed.
    /// Since the matrix is square, the `j`th row will also be removed.
    pub fn remove_column(
        &self,
        j: usize,
    ) -> Cholesky<N, DimDiff<D, U1>>
        where
            D: DimSub<U1>,
            DefaultAllocator: Allocator<N, DimDiff<D, U1>, DimDiff<D, U1>> + Allocator<N, D>
    {
        let n = self.chol.nrows();
        assert!(n > 0, "The matrix needs at least one column.");
        assert!(j < n, "j needs to be within the bound of the matrix.");

        // loads the data into a new matrix except for the jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.sub(U1), self.chol.data.shape().1.sub(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j..).copy_from(&self.chol.slice_range(..j, j + 1..));
        chol.slice_range_mut(j.., ..j).copy_from(&self.chol.slice_range(j + 1.., ..j));
        chol.slice_range_mut(j.., j..).copy_from(&self.chol.slice_range(j + 1.., j + 1..));

        // updates the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j.., j..);
        let mut workspace = self.chol.column(j).clone_owned();
        let mut old_colj = workspace.rows_range_mut(j + 1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut old_colj, N::RealField::one());

        Cholesky { chol }
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `x`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (x * x.adjoint())`.
    ///
    /// This helper method is called by `rank_one_update` but also `insert_column` and `remove_column`
    /// where it is used on a square slice of the decomposition
    fn xx_rank_one_update<Dm, Sm, Rx, Sx>(chol : &mut Matrix<N, Dm, Dm, Sm>, x: &mut Vector<N, Rx, Sx>, sigma: N::RealField)
        where
            //N: ComplexField,
            Dm: Dim,
            Rx: Dim,
            Sm: StorageMut<N, Dm, Dm>,
            Sx: StorageMut<N, Rx, U1>,
    {
        // heavily inspired by Eigen's `llt_rank_update_lower` implementation https://eigen.tuxfamily.org/dox/LLT_8h_source.html
        let n = x.nrows();
        assert_eq!(
            n,
            chol.nrows(),
            "The input vector must be of the same size as the factorized matrix."
        );

        let mut beta = crate::one::<N::RealField>();

        for j in 0..n {
            // updates the diagonal
            let diag = N::real(unsafe { *chol.get_unchecked((j, j)) });
            let diag2 = diag * diag;
            let xj = unsafe { *x.get_unchecked(j) };
            let sigma_xj2 = sigma * N::modulus_squared(xj);
            let gamma = diag2 * beta + sigma_xj2;
            let new_diag = (diag2 + sigma_xj2 / beta).sqrt();
            unsafe { *chol.get_unchecked_mut((j, j)) = N::from_real(new_diag) };
            beta += sigma_xj2 / diag2;
            // updates the terms of L
            let mut xjplus = x.rows_range_mut(j + 1..);
            let mut col_j = chol.slice_range_mut(j + 1.., j);
            // temp_jplus -= (wj / N::from_real(diag)) * col_j;
            xjplus.axpy(-xj / N::from_real(diag), &col_j, N::one());
            if gamma != crate::zero::<N::RealField>() {
                // col_j = N::from_real(nljj / diag) * col_j  + (N::from_real(nljj * sigma / gamma) * N::conjugate(wj)) * temp_jplus;
                col_j.axpy(
                    N::from_real(new_diag * sigma / gamma) * N::conjugate(xj),
                    &xjplus,
                    N::from_real(new_diag / diag),
                );
            }
        }
    }
}

impl<N: ComplexField, D: DimSub<Dynamic>, S: Storage<N, D, D>> SquareMatrix<N, D, S>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of this matrix.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn cholesky(self) -> Option<Cholesky<N, D>> {
        Cholesky::new(self.into_owned())
    }
}
<commit_msg>Unnecessary Trait bound DimSub<Dynamic> in impl Cholesky<commit_after>#[cfg(feature = "serde-serialize")]
use serde::{Deserialize, Serialize};

use num::One;
use alga::general::ComplexField;

use crate::allocator::Allocator;
use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, SquareMatrix, Vector};
use crate::constraint::{SameNumberOfRows, ShapeConstraint};
use crate::dimension::{Dim, DimAdd, DimSum, DimDiff, DimSub, Dynamic, U1};
use crate::storage::{Storage, StorageMut};

/// The Cholesky decomposition of a symmetric-definite-positive matrix.
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(serialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Serialize"))
)]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(deserialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Deserialize<'de>"))
)]
#[derive(Clone, Debug)]
pub struct Cholesky<N: ComplexField, D: Dim>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    chol: MatrixN<N, D>,
}

impl<N: ComplexField, D: Dim> Copy for Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
    MatrixN<N, D>: Copy,
{
}

impl<N: ComplexField, D: Dim> Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of `matrix`.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn new(mut matrix: MatrixN<N, D>) -> Option<Self> {
        assert!(matrix.is_square(), "The input matrix must be square.");

        let n = matrix.nrows();

        for j in 0..n {
            for k in 0..j {
                let factor = unsafe { -*matrix.get_unchecked((j, k)) };

                let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
                let mut col_j = col_j.rows_range_mut(j..);
                let col_k = col_k.rows_range(j..);

                col_j.axpy(factor.conjugate(), &col_k, N::one());
            }

            let diag = unsafe { *matrix.get_unchecked((j, j)) };
            if !diag.is_zero() {
                if let Some(denom) = diag.try_sqrt() {
                    unsafe {
                        *matrix.get_unchecked_mut((j, j)) = denom;
                    }

                    let mut col = matrix.slice_range_mut(j + 1.., j);
                    col /= denom;
                    continue;
                }
            }

            // The diagonal element is either zero or its square root could not
            // be taken (e.g. for negative real numbers).
            return None;
        }

        Some(Cholesky { chol: matrix })
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// upper-triangular part filled with zeros.
    pub fn unpack(mut self) -> MatrixN<N, D> {
        self.chol.fill_upper_triangle(N::zero(), 1);
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// The values of the strict upper-triangular part are garbage and should be ignored by further
    /// computations.
    pub fn unpack_dirty(self) -> MatrixN<N, D> {
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// uppen-triangular part filled with zeros.
    pub fn l(&self) -> MatrixN<N, D> {
        self.chol.lower_triangle()
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// This is an allocation-less version of `self.l()`. The values of the strict upper-triangular
    /// part are garbage and should be ignored by further computations.
    pub fn l_dirty(&self) -> &MatrixN<N, D> {
        &self.chol
    }

    /// Solves the system `self * x = b` where `self` is the decomposed matrix and `x` the unknown.
    ///
    /// The result is stored on `b`.
    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&self, b: &mut Matrix<N, R2, C2, S2>)
    where
        S2: StorageMut<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let _ = self.chol.solve_lower_triangular_mut(b);
        let _ = self.chol.ad_solve_lower_triangular_mut(b);
    }

    /// Returns the solution of the system `self * x = b` where `self` is the decomposed matrix and
    /// `x` the unknown.
    pub fn solve<R2: Dim, C2: Dim, S2>(&self, b: &Matrix<N, R2, C2, S2>) -> MatrixMN<N, R2, C2>
    where
        S2: Storage<N, R2, C2>,
        DefaultAllocator: Allocator<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let mut res = b.clone_owned();
        self.solve_mut(&mut res);
        res
    }

    /// Computes the inverse of the decomposed matrix.
    pub fn inverse(&self) -> MatrixN<N, D> {
        let shape = self.chol.data.shape();
        let mut res = MatrixN::identity_generic(shape.0, shape.1);

        self.solve_mut(&mut res);
        res
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (v * v.adjoint())`.
    #[inline]
    pub fn rank_one_update<R2: Dim, S2>(&mut self, x: &Vector<N, R2, S2>, sigma: N::RealField)
    where
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, R2, U1>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        Self::xx_rank_one_update(&mut self.chol, &mut x.clone_owned(), sigma)
    }

    /// Updates the decomposition such that we get the decomposition of a matrix with the given column `col` in the `j`th position.
    /// Since the matrix is square, an identical row will be added in the `j`th row.
    pub fn insert_column<R2, S2>(
        &self,
        j: usize,
        col: Vector<N, R2, S2>,
    ) -> Cholesky<N, DimSum<D, U1>>
    where
        D: DimAdd<U1>,
        R2: Dim,
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, DimSum<D, U1>, DimSum<D, U1>> + Allocator<N, R2>,
        ShapeConstraint: SameNumberOfRows<R2, DimSum<D, U1>>,
    {
        let mut col = col.into_owned();
        // for an explanation of the formulas, see https://en.wikipedia.org/wiki/Cholesky_decomposition#Updating_the_decomposition
        let n = col.nrows();
        assert_eq!(n, self.chol.nrows() + 1, "The new column must have the size of the factored matrix plus one.");
        assert!(j < n, "j needs to be within the bound of the new matrix.");

        // loads the data into a new matrix with an additional jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.add(U1), self.chol.data.shape().1.add(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j + 1..).copy_from(&self.chol.slice_range(..j, j..));
        chol.slice_range_mut(j + 1.., ..j).copy_from(&self.chol.slice_range(j.., ..j));
        chol.slice_range_mut(j + 1.., j + 1..).copy_from(&self.chol.slice_range(j.., j..));

        // update the jth row
        let top_left_corner = self.chol.slice_range(..j, ..j);

        let col_j = col[j];
        let (mut new_rowj_adjoint, mut new_colj) = col.rows_range_pair_mut(..j, j + 1..);
        assert!(top_left_corner.solve_lower_triangular_mut(&mut new_rowj_adjoint), "Cholesky::insert_column : Unable to solve lower triangular system!");

        new_rowj_adjoint.adjoint_to(&mut chol.slice_range_mut(j, ..j));

        // update the center element
        let center_element = N::sqrt(col_j - N::from_real(new_rowj_adjoint.norm_squared()));
        chol[(j, j)] = center_element;

        // update the jth column
        let bottom_left_corner = self.chol.slice_range(j.., ..j);
        // new_colj = (col_jplus - bottom_left_corner * new_rowj.adjoint()) / center_element;
        new_colj.gemm(-N::one() / center_element, &bottom_left_corner, &new_rowj_adjoint, N::one() / center_element);
        chol.slice_range_mut(j + 1.., j).copy_from(&new_colj);

        // update the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j + 1.., j + 1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut new_colj, -N::RealField::one());

        Cholesky { chol }
    }

    /// Updates the decomposition such that we get the decomposition of the factored matrix with its `j`th column removed.
    /// Since the matrix is square, the `j`th row will also be removed.
    pub fn remove_column(
        &self,
        j: usize,
    ) -> Cholesky<N, DimDiff<D, U1>>
        where
            D: DimSub<U1>,
            DefaultAllocator: Allocator<N, DimDiff<D, U1>, DimDiff<D, U1>> + Allocator<N, D>
    {
        let n = self.chol.nrows();
        assert!(n > 0, "The matrix needs at least one column.");
        assert!(j < n, "j needs to be within the bound of the matrix.");

        // loads the data into a new matrix except for the jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.sub(U1), self.chol.data.shape().1.sub(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j..).copy_from(&self.chol.slice_range(..j, j + 1..));
        chol.slice_range_mut(j.., ..j).copy_from(&self.chol.slice_range(j + 1.., ..j));
        chol.slice_range_mut(j.., j..).copy_from(&self.chol.slice_range(j + 1.., j + 1..));

        // updates the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j.., j..);
        let mut workspace = self.chol.column(j).clone_owned();
        let mut old_colj = workspace.rows_range_mut(j + 1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut old_colj, N::RealField::one());

        Cholesky { chol }
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `x`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (x * x.adjoint())`.
    ///
    /// This helper method is called by `rank_one_update` but also `insert_column` and `remove_column`
    /// where it is used on a square slice of the decomposition
    fn xx_rank_one_update<Dm, Sm, Rx, Sx>(chol : &mut Matrix<N, Dm, Dm, Sm>, x: &mut Vector<N, Rx, Sx>, sigma: N::RealField)
        where
            //N: ComplexField,
            Dm: Dim,
            Rx: Dim,
            Sm: StorageMut<N, Dm, Dm>,
            Sx: StorageMut<N, Rx, U1>,
    {
        // heavily inspired by Eigen's `llt_rank_update_lower` implementation https://eigen.tuxfamily.org/dox/LLT_8h_source.html
        let n = x.nrows();
        assert_eq!(
            n,
            chol.nrows(),
            "The input vector must be of the same size as the factorized matrix."
        );

        let mut beta = crate::one::<N::RealField>();

        for j in 0..n {
            // updates the diagonal
            let diag = N::real(unsafe { *chol.get_unchecked((j, j)) });
            let diag2 = diag * diag;
            let xj = unsafe { *x.get_unchecked(j) };
            let sigma_xj2 = sigma * N::modulus_squared(xj);
            let gamma = diag2 * beta + sigma_xj2;
            let new_diag = (diag2 + sigma_xj2 / beta).sqrt();
            unsafe { *chol.get_unchecked_mut((j, j)) = N::from_real(new_diag) };
            beta += sigma_xj2 / diag2;
            // updates the terms of L
            let mut xjplus = x.rows_range_mut(j + 1..);
            let mut col_j = chol.slice_range_mut(j + 1.., j);
            // temp_jplus -= (wj / N::from_real(diag)) * col_j;
            xjplus.axpy(-xj / N::from_real(diag), &col_j, N::one());
            if gamma != crate::zero::<N::RealField>() {
                // col_j = N::from_real(nljj / diag) * col_j  + (N::from_real(nljj * sigma / gamma) * N::conjugate(wj)) * temp_jplus;
                col_j.axpy(
                    N::from_real(new_diag * sigma / gamma) * N::conjugate(xj),
                    &xjplus,
                    N::from_real(new_diag / diag),
                );
            }
        }
    }
}

impl<N: ComplexField, D: DimSub<Dynamic>, S: Storage<N, D, D>> SquareMatrix<N, D, S>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of this matrix.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn cholesky(self) -> Option<Cholesky<N, D>> {
        Cholesky::new(self.into_owned())
    }
}
<|endoftext|>