<commit_before>use nom::be_u16;

/// Query operation types
#[derive(Debug,Clone,Copy,PartialEq)]
pub enum Opcode {
    Query,
    InverseQuery,
    Status,
}

/// Response types
#[derive(Debug,Clone,Copy,PartialEq)]
pub enum Rcode {
    NoError,
    FormatError,
    ServerFailure,
    NameError,
    NotImplemented,
    Refused,
}

/// Header for resource record queries and responses
#[derive(Debug,Clone,Copy,PartialEq)]
pub struct Header {
    /// A 16 bit identifier assigned by the program.
    pub id: u16,
    /// Specifies whether this message is a query (`false`) or response (`true`).
    pub qr: bool,
    /// The type of query
    pub opcode: Opcode,
    /// Whether the response is authoritative
    pub authoritative: bool,
    /// Whether the response is truncated
    pub truncated: bool,
    /// Whether recursion is desired
    pub recursion_desired: bool,
    /// Whether recursion is available
    pub recursion_available: bool,
    /// The response code
    pub rcode: Rcode,
    /// The number of entries in the question section.
    pub question_count: u16,
    /// The number of entries in the resource records section.
    pub answer_count: u16,
    /// The number of entries in the authority records section.
    pub ns_count: u16,
    /// The number of entries in the additional records section.
    pub additional_count: u16,
}

impl Header {
    /// Create a `Header` for a query
    pub fn query(id: u16, opcode: Opcode, recursion_desired: bool, questions: u16) -> Header {
        Header {
            id: id,
            qr: false,
            opcode: opcode,
            authoritative: false,
            truncated: false,
            recursion_desired: recursion_desired,
            recursion_available: false,
            rcode: Rcode::NoError,
            question_count: questions,
            answer_count: 0,
            ns_count: 0,
            additional_count: 0,
        }
    }

    /// Create a `Header` for a response
    pub fn response(query: Header, recursion_available: bool) -> Header {
        Header {
            id: query.id,
            qr: true,
            opcode: query.opcode,
            authoritative: false,
            truncated: false,
            recursion_desired: query.recursion_desired,
            recursion_available: recursion_available,
            rcode: Rcode::NoError,
            question_count: query.question_count,
            answer_count: 0,
            ns_count: 0,
            additional_count: 0,
        }
    }

    /// Creates a copy of the header, with the `answer_count` field modified.
    pub fn answers(&self, count: u16) -> Header {
        Header {
            id: self.id,
            qr: self.qr,
            opcode: self.opcode,
            authoritative: self.authoritative,
            truncated: self.truncated,
            recursion_desired: self.recursion_desired,
            recursion_available: self.recursion_available,
            rcode: self.rcode,
            question_count: self.question_count,
            answer_count: count,
            ns_count: self.ns_count,
            additional_count: self.additional_count,
        }
    }

    /// Creates a copy of the header, with the `ns_count` field modified.
    pub fn authorities(&self, count: u16) -> Header {
        Header {
            id: self.id,
            qr: self.qr,
            opcode: self.opcode,
            authoritative: self.authoritative,
            truncated: self.truncated,
            recursion_desired: self.recursion_desired,
            recursion_available: self.recursion_available,
            rcode: self.rcode,
            question_count: self.question_count,
            answer_count: self.answer_count,
            ns_count: count,
            additional_count: self.additional_count,
        }
    }

    /// Creates a copy of the header, with the `additional_count` field modified.
    pub fn additional(&self, count: u16) -> Header {
        Header {
            id: self.id,
            qr: self.qr,
            opcode: self.opcode,
            authoritative: self.authoritative,
            truncated: self.truncated,
            recursion_desired: self.recursion_desired,
            recursion_available: self.recursion_available,
            rcode: self.rcode,
            question_count: self.question_count,
            answer_count: self.answer_count,
            ns_count: self.ns_count,
            additional_count: count,
        }
    }
}

//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      ID                       |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |QR|   Opcode  |AA|TC|RD|RA|   Z    |   RCODE   |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    QDCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ANCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    NSCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ARCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

fn opcode_from(bits: u8) -> Option<Opcode> {
    match bits {
        0 => Some(Opcode::Query),
        1 => Some(Opcode::InverseQuery),
        2 => Some(Opcode::Status),
        _ => None,
    }
}
fn rcode_from(bits: u8) -> Option<Rcode> {
    match bits {
        0 => Some(Rcode::NoError),
        1 => Some(Rcode::FormatError),
        2 => Some(Rcode::ServerFailure),
        3 => Some(Rcode::NameError),
        4 => Some(Rcode::NotImplemented),
        5 => Some(Rcode::Refused),
        _ => None,
    }
}

named!(header_flags<&[u8], (bool, Opcode, bool, bool, bool, bool, Rcode)>,
bits!(do_parse!(
     qr:     take_bits!( u8, 1 ) >>
     opcode: map_opt!(take_bits!( u8, 4 ), opcode_from) >>
     aa:     take_bits!( u8, 1 ) >>
     tc:     take_bits!( u8, 1 ) >>
     rd:     take_bits!( u8, 1 ) >>
     ra:     take_bits!( u8, 1 ) >>
     zero:   take_bits!( u8, 3 ) >>
     rcode:  map_opt!(take_bits!( u8, 4 ), rcode_from) >>
     (((qr == 1), opcode, (aa == 1), (tc == 1), (rd == 1), (ra == 1), rcode))
)));

named!(pub parse_header<&[u8], Header>,
do_parse!(
    id:          be_u16 >>
    flags: header_flags >>
    qdcount:     be_u16 >>
    ancount:     be_u16 >>
    nscount:     be_u16 >>
    arcount:     be_u16 >>
    (Header {
        id: id,
        qr: flags.0,
        opcode: flags.1,
        authoritative: flags.2,
        truncated: flags.3,
        recursion_desired: flags.4,
        recursion_available: flags.5,
        rcode: flags.6,
        question_count: qdcount,
        answer_count: ancount,
        ns_count: nscount,
        additional_count: arcount
    })
));

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

    fn query_1() -> Header {
        Header::query(2, Opcode::Query, true, 1)
    }
    fn response_1() -> Header {
        Header::response(query_1(), true).answers(1)
    }

    #[test]
    fn parse_query_1_header() {
        let data = include_bytes!("../assets/captures/dns_1_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_1()));
    }

    #[test]
    fn parse_response_1_header() {
        let data = include_bytes!("../assets/captures/dns_1_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_1()));
    }

    fn query_2() -> Header {
        Header::query(3, Opcode::Query, true, 1)
    }
    fn response_2() -> Header {
        Header::response(query_2(), true).answers(1)
    }

    #[test]
    fn parse_query_2_header() {
        let data = include_bytes!("../assets/captures/dns_2_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_2()));
    }

    #[test]
    fn parse_response_2_header() {
        let data = include_bytes!("../assets/captures/dns_2_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_2()));
    }

    fn query_3() -> Header {
        Header::query(0xda64, Opcode::Query, true, 1)
    }
    fn response_3() -> Header {
        Header::response(query_3(), true)
            .answers(2)
            .authorities(1)
    }

    #[test]
    fn parse_query_3_header() {
        let data = include_bytes!("../assets/captures/dns_3_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_3()));
    }

    #[test]
    fn parse_response_3_header() {
        let data = include_bytes!("../assets/captures/dns_3_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_3()));
    }

    fn query_4() -> Header {
        Header::query(0x60ff, Opcode::Query, true, 1).additional(1)
    }
    fn response_4() -> Header {
        Header::response(query_4(), true)
            .answers(13)
            .additional(1)
    }

    #[test]
    fn parse_query_4_header() {
        let data = include_bytes!("../assets/captures/dns_4_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_4()));
    }

    #[test]
    fn parse_response_4_header() {
        let data = include_bytes!("../assets/captures/dns_4_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_4()));
    }
}
<commit_msg>feat: add 'unknown' values for opcodes<commit_after>use nom::be_u16;

/// Query operation types
#[derive(Debug,Clone,Copy,PartialEq)]
pub enum Opcode {
    Query,
    InverseQuery,
    Status,
    Unknown { value: u8 }
}

/// Response types
#[derive(Debug,Clone,Copy,PartialEq)]
pub enum Rcode {
    NoError,
    FormatError,
    ServerFailure,
    NameError,
    NotImplemented,
    Refused,
    Unknown { value: u8 },
}

/// Header for resource record queries and responses
#[derive(Debug,Clone,Copy,PartialEq)]
pub struct Header {
    /// A 16 bit identifier assigned by the program.
    pub id: u16,
    /// Specifies whether this message is a query (`false`) or response (`true`).
    pub qr: bool,
    /// The type of query
    pub opcode: Opcode,
    /// Whether the response is authoritative
    pub authoritative: bool,
    /// Whether the response is truncated
    pub truncated: bool,
    /// Whether recursion is desired
    pub recursion_desired: bool,
    /// Whether recursion is available
    pub recursion_available: bool,
    /// The response code
    pub rcode: Rcode,
    /// The number of entries in the question section.
    pub question_count: u16,
    /// The number of entries in the resource records section.
    pub answer_count: u16,
    /// The number of entries in the authority records section.
    pub ns_count: u16,
    /// The number of entries in the additional records section.
    pub additional_count: u16,
}

impl Header {
    /// Create a `Header` for a query
    pub fn query(id: u16, opcode: Opcode, recursion_desired: bool, questions: u16) -> Header {
        Header {
            id: id,
            qr: false,
            opcode: opcode,
            authoritative: false,
            truncated: false,
            recursion_desired: recursion_desired,
            recursion_available: false,
            rcode: Rcode::NoError,
            question_count: questions,
            answer_count: 0,
            ns_count: 0,
            additional_count: 0,
        }
    }

    /// Create a `Header` for a response
    pub fn response(query: Header, recursion_available: bool) -> Header {
        Header {
            id: query.id,
            qr: true,
            opcode: query.opcode,
            authoritative: false,
            truncated: false,
            recursion_desired: query.recursion_desired,
            recursion_available: recursion_available,
            rcode: Rcode::NoError,
            question_count: query.question_count,
            answer_count: 0,
            ns_count: 0,
            additional_count: 0,
        }
    }

    /// Creates a copy of the header, with the `answer_count` field modified.
    pub fn answers(&self, count: u16) -> Header {
        Header {
            id: self.id,
            qr: self.qr,
            opcode: self.opcode,
            authoritative: self.authoritative,
            truncated: self.truncated,
            recursion_desired: self.recursion_desired,
            recursion_available: self.recursion_available,
            rcode: self.rcode,
            question_count: self.question_count,
            answer_count: count,
            ns_count: self.ns_count,
            additional_count: self.additional_count,
        }
    }

    /// Creates a copy of the header, with the `ns_count` field modified.
    pub fn authorities(&self, count: u16) -> Header {
        Header {
            id: self.id,
            qr: self.qr,
            opcode: self.opcode,
            authoritative: self.authoritative,
            truncated: self.truncated,
            recursion_desired: self.recursion_desired,
            recursion_available: self.recursion_available,
            rcode: self.rcode,
            question_count: self.question_count,
            answer_count: self.answer_count,
            ns_count: count,
            additional_count: self.additional_count,
        }
    }

    /// Creates a copy of the header, with the `additional_count` field modified.
    pub fn additional(&self, count: u16) -> Header {
        Header {
            id: self.id,
            qr: self.qr,
            opcode: self.opcode,
            authoritative: self.authoritative,
            truncated: self.truncated,
            recursion_desired: self.recursion_desired,
            recursion_available: self.recursion_available,
            rcode: self.rcode,
            question_count: self.question_count,
            answer_count: self.answer_count,
            ns_count: self.ns_count,
            additional_count: count,
        }
    }
}

//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      ID                       |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |QR|   Opcode  |AA|TC|RD|RA|   Z    |   RCODE   |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    QDCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ANCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    NSCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ARCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

fn opcode_from(bits: u8) -> Opcode {
    match bits {
        0 => Opcode::Query,
        1 => Opcode::InverseQuery,
        2 => Opcode::Status,
        x @ _ => Opcode::Unknown { value: x },
    }
}
fn rcode_from(bits: u8) -> Rcode {
    match bits {
        0 => Rcode::NoError,
        1 => Rcode::FormatError,
        2 => Rcode::ServerFailure,
        3 => Rcode::NameError,
        4 => Rcode::NotImplemented,
        5 => Rcode::Refused,
        x @ _ => Rcode::Unknown { value: x },
    }
}

named!(header_flags<&[u8], (bool, Opcode, bool, bool, bool, bool, Rcode)>,
bits!(do_parse!(
     qr:     take_bits!( u8, 1 ) >>
     opcode: map!(take_bits!( u8, 4 ), opcode_from) >>
     aa:     take_bits!( u8, 1 ) >>
     tc:     take_bits!( u8, 1 ) >>
     rd:     take_bits!( u8, 1 ) >>
     ra:     take_bits!( u8, 1 ) >>
     zero:   take_bits!( u8, 3 ) >>
     rcode:  map!(take_bits!( u8, 4 ), rcode_from) >>
     (((qr == 1), opcode, (aa == 1), (tc == 1), (rd == 1), (ra == 1), rcode))
)));

named!(pub parse_header<&[u8], Header>,
do_parse!(
    id:          be_u16 >>
    flags: header_flags >>
    qdcount:     be_u16 >>
    ancount:     be_u16 >>
    nscount:     be_u16 >>
    arcount:     be_u16 >>
    (Header {
        id: id,
        qr: flags.0,
        opcode: flags.1,
        authoritative: flags.2,
        truncated: flags.3,
        recursion_desired: flags.4,
        recursion_available: flags.5,
        rcode: flags.6,
        question_count: qdcount,
        answer_count: ancount,
        ns_count: nscount,
        additional_count: arcount
    })
));

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

    fn query_1() -> Header {
        Header::query(2, Opcode::Query, true, 1)
    }
    fn response_1() -> Header {
        Header::response(query_1(), true).answers(1)
    }

    #[test]
    fn parse_query_1_header() {
        let data = include_bytes!("../assets/captures/dns_1_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_1()));
    }

    #[test]
    fn parse_response_1_header() {
        let data = include_bytes!("../assets/captures/dns_1_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_1()));
    }

    fn query_2() -> Header {
        Header::query(3, Opcode::Query, true, 1)
    }
    fn response_2() -> Header {
        Header::response(query_2(), true).answers(1)
    }

    #[test]
    fn parse_query_2_header() {
        let data = include_bytes!("../assets/captures/dns_2_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_2()));
    }

    #[test]
    fn parse_response_2_header() {
        let data = include_bytes!("../assets/captures/dns_2_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_2()));
    }

    fn query_3() -> Header {
        Header::query(0xda64, Opcode::Query, true, 1)
    }
    fn response_3() -> Header {
        Header::response(query_3(), true)
            .answers(2)
            .authorities(1)
    }

    #[test]
    fn parse_query_3_header() {
        let data = include_bytes!("../assets/captures/dns_3_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_3()));
    }

    #[test]
    fn parse_response_3_header() {
        let data = include_bytes!("../assets/captures/dns_3_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_3()));
    }

    fn query_4() -> Header {
        Header::query(0x60ff, Opcode::Query, true, 1).additional(1)
    }
    fn response_4() -> Header {
        Header::response(query_4(), true)
            .answers(13)
            .additional(1)
    }

    #[test]
    fn parse_query_4_header() {
        let data = include_bytes!("../assets/captures/dns_4_query.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], query_4()));
    }

    #[test]
    fn parse_response_4_header() {
        let data = include_bytes!("../assets/captures/dns_4_response.bin");
        assert_eq!(parse_header(&data[0..12]), Done(&b""[..], response_4()));
    }
}
<|endoftext|>


<commit_before>//! Assertions for array

use ndarray::*;

use super::norm::*;
use super::types::*;

/// check two values are close in terms of the relative torrence
pub fn rclose<A: Scalar>(test: A, truth: A, rtol: A::Real) -> Result<A::Real, A::Real> {
    let dev = (test - truth).abs() / truth.abs();
    if dev < rtol {
        Ok(dev)
    } else {
        Err(dev)
    }
}

/// check two values are close in terms of the absolute torrence
pub fn aclose<A: Scalar>(test: A, truth: A, atol: A::Real) -> Result<A::Real, A::Real> {
    let dev = (test - truth).abs();
    if dev < atol {
        Ok(dev)
    } else {
        Err(dev)
    }
}

/// check two arrays are close in maximum norm
pub fn close_max<A, S1, S2, D>(
    test: &ArrayBase<S1, D>,
    truth: &ArrayBase<S2, D>,
    atol: A::Real,
) -> Result<A::Real, A::Real>
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: Data<Elem = A>,
    D: Dimension,
{
    let tol = (test - truth).norm_max();
    if tol < atol {
        Ok(tol)
    } else {
        Err(tol)
    }
}

/// check two arrays are close in L1 norm
pub fn close_l1<A, S1, S2, D>(
    test: &ArrayBase<S1, D>,
    truth: &ArrayBase<S2, D>,
    rtol: A::Real,
) -> Result<A::Real, A::Real>
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: Data<Elem = A>,
    D: Dimension,
{
    let tol = (test - truth).norm_l1() / truth.norm_l1();
    if tol < rtol {
        Ok(tol)
    } else {
        Err(tol)
    }
}

/// check two arrays are close in L2 norm
pub fn close_l2<A, S1, S2, D>(
    test: &ArrayBase<S1, D>,
    truth: &ArrayBase<S2, D>,
    rtol: A::Real,
) -> Result<A::Real, A::Real>
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: Data<Elem = A>,
    D: Dimension,
{
    let tol = (test - truth).norm_l2() / truth.norm_l2();
    if tol < rtol {
        Ok(tol)
    } else {
        Err(tol)
    }
}

macro_rules! generate_assert {
    ($assert:ident, $close:path) => {
        #[macro_export]
        macro_rules! $assert {
            ($test: expr,$truth: expr,$tol: expr) => {
                $crate::$close($test, $truth, $tol).unwrap();
            };
        }
    };
} // generate_assert!

generate_assert!(assert_rclose, rclose);
generate_assert!(assert_aclose, aclose);
generate_assert!(assert_close_max, close_max);
generate_assert!(assert_close_l1, close_l1);
generate_assert!(assert_close_l2, close_l2);
<commit_msg>Panic from close_* functions (not as unwrap)<commit_after>//! Assertions for array

use ndarray::*;
use std::fmt::Debug;

use super::norm::*;
use super::types::*;

/// check two values are close in terms of the relative torrence
pub fn rclose<A: Scalar>(test: A, truth: A, rtol: A::Real) {
    let dev = (test - truth).abs() / truth.abs();
    if dev > rtol {
        eprintln!("Expected = {}", truth);
        eprintln!("Actual   = {}", test);
        panic!("Too large deviation in relative torrence: {}", dev);
    }
}

/// check two values are close in terms of the absolute torrence
pub fn aclose<A: Scalar>(test: A, truth: A, atol: A::Real) {
    let dev = (test - truth).abs();
    if dev > atol {
        eprintln!("Expected = {}", truth);
        eprintln!("Actual   = {}", test);
        panic!("Too large deviation in absolute torrence: {}", dev);
    }
}

/// check two arrays are close in maximum norm
pub fn close_max<A, S1, S2, D>(test: &ArrayBase<S1, D>, truth: &ArrayBase<S2, D>, atol: A::Real)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: Data<Elem = A>,
    D: Dimension,
    D::Pattern: PartialEq + Debug,
{
    let tol = (test - truth).norm_max();
    if tol > atol {
        eprintln!("Expected:\n{}", truth);
        eprintln!("Actual:\n{}", test);
        panic!("Too large deviation in maximum norm: {}", tol);
    }
}

/// check two arrays are close in L1 norm
pub fn close_l1<A, S1, S2, D>(test: &ArrayBase<S1, D>, truth: &ArrayBase<S2, D>, rtol: A::Real)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: Data<Elem = A>,
    D: Dimension,
    D::Pattern: PartialEq + Debug,
{
    let tol = (test - truth).norm_l1() / truth.norm_l1();
    if tol > rtol {
        eprintln!("Expected:\n{}", truth);
        eprintln!("Actual:\n{}", test);
        panic!("Too large deviation in L1-norm: {}", tol);
    }
}

/// check two arrays are close in L2 norm
pub fn close_l2<A, S1, S2, D>(test: &ArrayBase<S1, D>, truth: &ArrayBase<S2, D>, rtol: A::Real)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: Data<Elem = A>,
    D: Dimension,
    D::Pattern: PartialEq + Debug,
{
    let tol = (test - truth).norm_l2() / truth.norm_l2();
    if tol > rtol {
        eprintln!("Expected:\n{}", truth);
        eprintln!("Actual:\n{}", test);
        panic!("Too large deviation in L2-norm: {}", tol);
    }
}

macro_rules! generate_assert {
    ($assert:ident, $close:path) => {
        #[macro_export]
        macro_rules! $assert {
            ($test: expr,$truth: expr,$tol: expr) => {
                $crate::$close($test, $truth, $tol);
            };
        }
    };
} // generate_assert!

generate_assert!(assert_rclose, rclose);
generate_assert!(assert_aclose, aclose);
generate_assert!(assert_close_max, close_max);
generate_assert!(assert_close_l1, close_l1);
generate_assert!(assert_close_l2, close_l2);
<|endoftext|>


<commit_before>//! Intermediate representation of variables.

use super::super::codegen::MacroTypeVariation;
use super::context::{BindgenContext, TypeId};
use super::dot::DotAttributes;
use super::function::cursor_mangling;
use super::int::IntKind;
use super::item::Item;
use super::ty::{FloatKind, TypeKind};
use crate::callbacks::MacroParsingBehavior;
use crate::clang;
use crate::clang::ClangToken;
use crate::parse::{
    ClangItemParser, ClangSubItemParser, ParseError, ParseResult,
};
use cexpr;
use std::io;
use std::num::Wrapping;

/// The type for a constant variable.
#[derive(Debug)]
pub enum VarType {
    /// A boolean.
    Bool(bool),
    /// An integer.
    Int(i64),
    /// A floating point number.
    Float(f64),
    /// A character.
    Char(u8),
    /// A string, not necessarily well-formed utf-8.
    String(Vec<u8>),
}

/// A `Var` is our intermediate representation of a variable.
#[derive(Debug)]
pub struct Var {
    /// The name of the variable.
    name: String,
    /// The mangled name of the variable.
    mangled_name: Option<String>,
    /// The type of the variable.
    ty: TypeId,
    /// The value of the variable, that needs to be suitable for `ty`.
    val: Option<VarType>,
    /// Whether this variable is const.
    is_const: bool,
}

impl Var {
    /// Construct a new `Var`.
    pub fn new(
        name: String,
        mangled_name: Option<String>,
        ty: TypeId,
        val: Option<VarType>,
        is_const: bool,
    ) -> Var {
        assert!(!name.is_empty());
        Var {
            name,
            mangled_name,
            ty,
            val,
            is_const,
        }
    }

    /// Is this variable `const` qualified?
    pub fn is_const(&self) -> bool {
        self.is_const
    }

    /// The value of this constant variable, if any.
    pub fn val(&self) -> Option<&VarType> {
        self.val.as_ref()
    }

    /// Get this variable's type.
    pub fn ty(&self) -> TypeId {
        self.ty
    }

    /// Get this variable's name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get this variable's mangled name.
    pub fn mangled_name(&self) -> Option<&str> {
        self.mangled_name.as_deref()
    }
}

impl DotAttributes for Var {
    fn dot_attributes<W>(
        &self,
        _ctx: &BindgenContext,
        out: &mut W,
    ) -> io::Result<()>
    where
        W: io::Write,
    {
        if self.is_const {
            writeln!(out, "
const	true
")?;
        }

        if let Some(ref mangled) = self.mangled_name {
            writeln!(
                out,
                "
mangled name	{}
",
                mangled
            )?;
        }

        Ok(())
    }
}

fn default_macro_constant_type(ctx: &BindgenContext, value: i64) -> IntKind {
    if value < 0 ||
        ctx.options().default_macro_constant_type ==
            MacroTypeVariation::Signed
    {
        if value < i32::min_value() as i64 || value > i32::max_value() as i64 {
            IntKind::I64
        } else if !ctx.options().fit_macro_constants ||
            value < i16::min_value() as i64 ||
            value > i16::max_value() as i64
        {
            IntKind::I32
        } else if value < i8::min_value() as i64 ||
            value > i8::max_value() as i64
        {
            IntKind::I16
        } else {
            IntKind::I8
        }
    } else if value > u32::max_value() as i64 {
        IntKind::U64
    } else if !ctx.options().fit_macro_constants ||
        value > u16::max_value() as i64
    {
        IntKind::U32
    } else if value > u8::max_value() as i64 {
        IntKind::U16
    } else {
        IntKind::U8
    }
}

/// Determines whether a set of tokens from a CXCursor_MacroDefinition
/// represent a function-like macro. If so, calls the func_macro callback
/// and returns `Err(ParseError::Continue)` to signal to skip further
/// processing. If conversion to UTF-8 fails (it is performed only where it
/// should be infallible), then `Err(ParseError::Continue)` is returned as well.
fn handle_function_macro(
    cursor: &clang::Cursor,
    tokens: &[ClangToken],
    callbacks: &dyn crate::callbacks::ParseCallbacks,
) -> Result<(), ParseError> {
    // TODO: Hoist the `is_macro_function_like` check into this function's
    // caller, and thus avoid allocating the `tokens` vector for non-functional
    // macros.
    let is_functional_macro = cursor.is_macro_function_like();

    if !is_functional_macro {
        return Ok(());
    }

    let is_closing_paren = |t: &ClangToken| {
        // Test cheap token kind before comparing exact spellings.
        t.kind == clang_sys::CXToken_Punctuation && t.spelling() == b")"
    };
    let boundary = tokens.iter().position(is_closing_paren);

    let mut spelled = tokens.iter().map(ClangToken::spelling);
    // Add 1, to convert index to length.
    let left = spelled
        .by_ref()
        .take(boundary.ok_or(ParseError::Continue)? + 1);
    let left = left.collect::<Vec<_>>().concat();
    let left = String::from_utf8(left).map_err(|_| ParseError::Continue)?;
    let right = spelled;
    // Drop last token with LLVM < 4.0, due to an LLVM bug.
    //
    // See:
    //   https://bugs.llvm.org//show_bug.cgi?id=9069
    let len = match (right.len(), crate::clang_version().parsed) {
        (len, Some((v, _))) if len > 0 && v < 4 => len - 1,
        (len, _) => len,
    };
    let right: Vec<_> = right.take(len).collect();
    callbacks.func_macro(&left, &right);

    // We handled the macro, skip future macro processing.
    Err(ParseError::Continue)
}

impl ClangSubItemParser for Var {
    fn parse(
        cursor: clang::Cursor,
        ctx: &mut BindgenContext,
    ) -> Result<ParseResult<Self>, ParseError> {
        use cexpr::expr::EvalResult;
        use cexpr::literal::CChar;
        use clang_sys::*;
        match cursor.kind() {
            CXCursor_MacroDefinition => {
                let tokens: Vec<_> = cursor.tokens().iter().collect();

                if let Some(callbacks) = ctx.parse_callbacks() {
                    match callbacks.will_parse_macro(&cursor.spelling()) {
                        MacroParsingBehavior::Ignore => {
                            return Err(ParseError::Continue);
                        }
                        MacroParsingBehavior::Default => {}
                    }

                    handle_function_macro(&cursor, &tokens, callbacks)?;
                }

                let value = parse_macro(ctx, &tokens);

                let (id, value) = match value {
                    Some(v) => v,
                    None => return Err(ParseError::Continue),
                };

                assert!(!id.is_empty(), "Empty macro name?");

                let previously_defined = ctx.parsed_macro(&id);

                // NB: It's important to "note" the macro even if the result is
                // not an integer, otherwise we might loose other kind of
                // derived macros.
                ctx.note_parsed_macro(id.clone(), value.clone());

                if previously_defined {
                    let name = String::from_utf8(id).unwrap();
                    warn!("Duplicated macro definition: {}", name);
                    return Err(ParseError::Continue);
                }

                // NOTE: Unwrapping, here and above, is safe, because the
                // identifier of a token comes straight from clang, and we
                // enforce utf8 there, so we should have already panicked at
                // this point.
                let name = String::from_utf8(id).unwrap();
                let (type_kind, val) = match value {
                    EvalResult::Invalid => return Err(ParseError::Continue),
                    EvalResult::Float(f) => {
                        (TypeKind::Float(FloatKind::Double), VarType::Float(f))
                    }
                    EvalResult::Char(c) => {
                        let c = match c {
                            CChar::Char(c) => {
                                assert_eq!(c.len_utf8(), 1);
                                c as u8
                            }
                            CChar::Raw(c) => {
                                assert!(c <= ::std::u8::MAX as u64);
                                c as u8
                            }
                        };

                        (TypeKind::Int(IntKind::U8), VarType::Char(c))
                    }
                    EvalResult::Str(val) => {
                        let char_ty = Item::builtin_type(
                            TypeKind::Int(IntKind::U8),
                            true,
                            ctx,
                        );
                        if let Some(callbacks) = ctx.parse_callbacks() {
                            callbacks.str_macro(&name, &val);
                        }
                        (TypeKind::Pointer(char_ty), VarType::String(val))
                    }
                    EvalResult::Int(Wrapping(value)) => {
                        let kind = ctx
                            .parse_callbacks()
                            .and_then(|c| c.int_macro(&name, value))
                            .unwrap_or_else(|| {
                                default_macro_constant_type(ctx, value)
                            });

                        (TypeKind::Int(kind), VarType::Int(value))
                    }
                };

                let ty = Item::builtin_type(type_kind, true, ctx);

                Ok(ParseResult::New(
                    Var::new(name, None, ty, Some(val), true),
                    Some(cursor),
                ))
            }
            CXCursor_VarDecl => {
                let name = cursor.spelling();
                if name.is_empty() {
                    warn!("Empty constant name?");
                    return Err(ParseError::Continue);
                }

                let ty = cursor.cur_type();

                // TODO(emilio): do we have to special-case constant arrays in
                // some other places?
                let is_const = ty.is_const() ||
                    (ty.kind() == CXType_ConstantArray &&
                        ty.elem_type()
                            .map_or(false, |element| element.is_const()));

                let ty = match Item::from_ty(&ty, cursor, None, ctx) {
                    Ok(ty) => ty,
                    Err(e) => {
                        assert_eq!(
                            ty.kind(),
                            CXType_Auto,
                            "Couldn't resolve constant type, and it \
                             wasn't an nondeductible auto type!"
                        );
                        return Err(e);
                    }
                };

                // Note: Ty might not be totally resolved yet, see
                // tests/headers/inner_const.hpp
                //
                // That's fine because in that case we know it's not a literal.
                let canonical_ty = ctx
                    .safe_resolve_type(ty)
                    .and_then(|t| t.safe_canonical_type(ctx));

                let is_integer = canonical_ty.map_or(false, |t| t.is_integer());
                let is_float = canonical_ty.map_or(false, |t| t.is_float());

                // TODO: We could handle `char` more gracefully.
                // TODO: Strings, though the lookup is a bit more hard (we need
                // to look at the canonical type of the pointee too, and check
                // is char, u8, or i8 I guess).
                let value = if is_integer {
                    let kind = match *canonical_ty.unwrap().kind() {
                        TypeKind::Int(kind) => kind,
                        _ => unreachable!(),
                    };

                    let mut val = cursor.evaluate().and_then(|v| v.as_int());
                    if val.is_none() || !kind.signedness_matches(val.unwrap()) {
                        let tu = ctx.translation_unit();
                        val = get_integer_literal_from_cursor(&cursor, tu);
                    }

                    val.map(|val| {
                        if kind == IntKind::Bool {
                            VarType::Bool(val != 0)
                        } else {
                            VarType::Int(val)
                        }
                    })
                } else if is_float {
                    cursor
                        .evaluate()
                        .and_then(|v| v.as_double())
                        .map(VarType::Float)
                } else {
                    cursor
                        .evaluate()
                        .and_then(|v| v.as_literal_string())
                        .map(VarType::String)
                };

                let mangling = cursor_mangling(ctx, &cursor);
                let var = Var::new(name, mangling, ty, value, is_const);

                Ok(ParseResult::New(var, Some(cursor)))
            }
            _ => {
                /* TODO */
                Err(ParseError::Continue)
            }
        }
    }
}

/// Try and parse a macro using all the macros parsed until now.
fn parse_macro(
    ctx: &BindgenContext,
    tokens: &[ClangToken],
) -> Option<(Vec<u8>, cexpr::expr::EvalResult)> {
    use cexpr::expr;

    let cexpr_tokens: Vec<_> = tokens
        .iter()
        .filter_map(ClangToken::as_cexpr_token)
        .collect();

    let parser = expr::IdentifierParser::new(ctx.parsed_macros());

    match parser.macro_definition(&cexpr_tokens) {
        Ok((_, (id, val))) => Some((id.into(), val)),
        _ => None,
    }
}

fn parse_int_literal_tokens(cursor: &clang::Cursor) -> Option<i64> {
    use cexpr::expr;
    use cexpr::expr::EvalResult;

    let cexpr_tokens = cursor.cexpr_tokens();

    // TODO(emilio): We can try to parse other kinds of literals.
    match expr::expr(&cexpr_tokens) {
        Ok((_, EvalResult::Int(Wrapping(val)))) => Some(val),
        _ => None,
    }
}

fn get_integer_literal_from_cursor(
    cursor: &clang::Cursor,
    unit: &clang::TranslationUnit,
) -> Option<i64> {
    use clang_sys::*;
    let mut value = None;
    cursor.visit(|c| {
        match c.kind() {
            CXCursor_IntegerLiteral | CXCursor_UnaryOperator => {
                value = parse_int_literal_tokens(&c);
            }
            CXCursor_UnexposedExpr => {
                value = get_integer_literal_from_cursor(&c, unit);
            }
            _ => (),
        }
        if value.is_some() {
            CXChildVisit_Break
        } else {
            CXChildVisit_Continue
        }
    });
    value
}
<commit_msg>parser: Simplify handle_function_macro for clang 4.0+<commit_after>//! Intermediate representation of variables.

use super::super::codegen::MacroTypeVariation;
use super::context::{BindgenContext, TypeId};
use super::dot::DotAttributes;
use super::function::cursor_mangling;
use super::int::IntKind;
use super::item::Item;
use super::ty::{FloatKind, TypeKind};
use crate::callbacks::MacroParsingBehavior;
use crate::clang;
use crate::clang::ClangToken;
use crate::parse::{
    ClangItemParser, ClangSubItemParser, ParseError, ParseResult,
};
use cexpr;
use std::io;
use std::num::Wrapping;

/// The type for a constant variable.
#[derive(Debug)]
pub enum VarType {
    /// A boolean.
    Bool(bool),
    /// An integer.
    Int(i64),
    /// A floating point number.
    Float(f64),
    /// A character.
    Char(u8),
    /// A string, not necessarily well-formed utf-8.
    String(Vec<u8>),
}

/// A `Var` is our intermediate representation of a variable.
#[derive(Debug)]
pub struct Var {
    /// The name of the variable.
    name: String,
    /// The mangled name of the variable.
    mangled_name: Option<String>,
    /// The type of the variable.
    ty: TypeId,
    /// The value of the variable, that needs to be suitable for `ty`.
    val: Option<VarType>,
    /// Whether this variable is const.
    is_const: bool,
}

impl Var {
    /// Construct a new `Var`.
    pub fn new(
        name: String,
        mangled_name: Option<String>,
        ty: TypeId,
        val: Option<VarType>,
        is_const: bool,
    ) -> Var {
        assert!(!name.is_empty());
        Var {
            name,
            mangled_name,
            ty,
            val,
            is_const,
        }
    }

    /// Is this variable `const` qualified?
    pub fn is_const(&self) -> bool {
        self.is_const
    }

    /// The value of this constant variable, if any.
    pub fn val(&self) -> Option<&VarType> {
        self.val.as_ref()
    }

    /// Get this variable's type.
    pub fn ty(&self) -> TypeId {
        self.ty
    }

    /// Get this variable's name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get this variable's mangled name.
    pub fn mangled_name(&self) -> Option<&str> {
        self.mangled_name.as_deref()
    }
}

impl DotAttributes for Var {
    fn dot_attributes<W>(
        &self,
        _ctx: &BindgenContext,
        out: &mut W,
    ) -> io::Result<()>
    where
        W: io::Write,
    {
        if self.is_const {
            writeln!(out, "
const	true
")?;
        }

        if let Some(ref mangled) = self.mangled_name {
            writeln!(
                out,
                "
mangled name	{}
",
                mangled
            )?;
        }

        Ok(())
    }
}

fn default_macro_constant_type(ctx: &BindgenContext, value: i64) -> IntKind {
    if value < 0 ||
        ctx.options().default_macro_constant_type ==
            MacroTypeVariation::Signed
    {
        if value < i32::min_value() as i64 || value > i32::max_value() as i64 {
            IntKind::I64
        } else if !ctx.options().fit_macro_constants ||
            value < i16::min_value() as i64 ||
            value > i16::max_value() as i64
        {
            IntKind::I32
        } else if value < i8::min_value() as i64 ||
            value > i8::max_value() as i64
        {
            IntKind::I16
        } else {
            IntKind::I8
        }
    } else if value > u32::max_value() as i64 {
        IntKind::U64
    } else if !ctx.options().fit_macro_constants ||
        value > u16::max_value() as i64
    {
        IntKind::U32
    } else if value > u8::max_value() as i64 {
        IntKind::U16
    } else {
        IntKind::U8
    }
}

/// Parses tokens from a CXCursor_MacroDefinition pointing into a function-like
/// macro, and calls the func_macro callback.
fn handle_function_macro(
    cursor: &clang::Cursor,
    callbacks: &dyn crate::callbacks::ParseCallbacks,
) {
    let is_closing_paren = |t: &ClangToken| {
        // Test cheap token kind before comparing exact spellings.
        t.kind == clang_sys::CXToken_Punctuation && t.spelling() == b")"
    };
    let tokens: Vec<_> = cursor.tokens().iter().collect();
    if let Some(boundary) = tokens.iter().position(is_closing_paren) {
        let mut spelled = tokens.iter().map(ClangToken::spelling);
        // Add 1, to convert index to length.
        let left = spelled.by_ref().take(boundary + 1);
        let left = left.collect::<Vec<_>>().concat();
        if let Ok(left) = String::from_utf8(left) {
            let right: Vec<_> = spelled.collect();
            callbacks.func_macro(&left, &right);
        }
    }
}

impl ClangSubItemParser for Var {
    fn parse(
        cursor: clang::Cursor,
        ctx: &mut BindgenContext,
    ) -> Result<ParseResult<Self>, ParseError> {
        use cexpr::expr::EvalResult;
        use cexpr::literal::CChar;
        use clang_sys::*;
        match cursor.kind() {
            CXCursor_MacroDefinition => {
                if let Some(callbacks) = ctx.parse_callbacks() {
                    match callbacks.will_parse_macro(&cursor.spelling()) {
                        MacroParsingBehavior::Ignore => {
                            return Err(ParseError::Continue);
                        }
                        MacroParsingBehavior::Default => {}
                    }

                    if cursor.is_macro_function_like() {
                        handle_function_macro(&cursor, callbacks);
                        // We handled the macro, skip macro processing below.
                        return Err(ParseError::Continue);
                    }
                }

                let value = parse_macro(ctx, &cursor);

                let (id, value) = match value {
                    Some(v) => v,
                    None => return Err(ParseError::Continue),
                };

                assert!(!id.is_empty(), "Empty macro name?");

                let previously_defined = ctx.parsed_macro(&id);

                // NB: It's important to "note" the macro even if the result is
                // not an integer, otherwise we might loose other kind of
                // derived macros.
                ctx.note_parsed_macro(id.clone(), value.clone());

                if previously_defined {
                    let name = String::from_utf8(id).unwrap();
                    warn!("Duplicated macro definition: {}", name);
                    return Err(ParseError::Continue);
                }

                // NOTE: Unwrapping, here and above, is safe, because the
                // identifier of a token comes straight from clang, and we
                // enforce utf8 there, so we should have already panicked at
                // this point.
                let name = String::from_utf8(id).unwrap();
                let (type_kind, val) = match value {
                    EvalResult::Invalid => return Err(ParseError::Continue),
                    EvalResult::Float(f) => {
                        (TypeKind::Float(FloatKind::Double), VarType::Float(f))
                    }
                    EvalResult::Char(c) => {
                        let c = match c {
                            CChar::Char(c) => {
                                assert_eq!(c.len_utf8(), 1);
                                c as u8
                            }
                            CChar::Raw(c) => {
                                assert!(c <= ::std::u8::MAX as u64);
                                c as u8
                            }
                        };

                        (TypeKind::Int(IntKind::U8), VarType::Char(c))
                    }
                    EvalResult::Str(val) => {
                        let char_ty = Item::builtin_type(
                            TypeKind::Int(IntKind::U8),
                            true,
                            ctx,
                        );
                        if let Some(callbacks) = ctx.parse_callbacks() {
                            callbacks.str_macro(&name, &val);
                        }
                        (TypeKind::Pointer(char_ty), VarType::String(val))
                    }
                    EvalResult::Int(Wrapping(value)) => {
                        let kind = ctx
                            .parse_callbacks()
                            .and_then(|c| c.int_macro(&name, value))
                            .unwrap_or_else(|| {
                                default_macro_constant_type(ctx, value)
                            });

                        (TypeKind::Int(kind), VarType::Int(value))
                    }
                };

                let ty = Item::builtin_type(type_kind, true, ctx);

                Ok(ParseResult::New(
                    Var::new(name, None, ty, Some(val), true),
                    Some(cursor),
                ))
            }
            CXCursor_VarDecl => {
                let name = cursor.spelling();
                if name.is_empty() {
                    warn!("Empty constant name?");
                    return Err(ParseError::Continue);
                }

                let ty = cursor.cur_type();

                // TODO(emilio): do we have to special-case constant arrays in
                // some other places?
                let is_const = ty.is_const() ||
                    (ty.kind() == CXType_ConstantArray &&
                        ty.elem_type()
                            .map_or(false, |element| element.is_const()));

                let ty = match Item::from_ty(&ty, cursor, None, ctx) {
                    Ok(ty) => ty,
                    Err(e) => {
                        assert_eq!(
                            ty.kind(),
                            CXType_Auto,
                            "Couldn't resolve constant type, and it \
                             wasn't an nondeductible auto type!"
                        );
                        return Err(e);
                    }
                };

                // Note: Ty might not be totally resolved yet, see
                // tests/headers/inner_const.hpp
                //
                // That's fine because in that case we know it's not a literal.
                let canonical_ty = ctx
                    .safe_resolve_type(ty)
                    .and_then(|t| t.safe_canonical_type(ctx));

                let is_integer = canonical_ty.map_or(false, |t| t.is_integer());
                let is_float = canonical_ty.map_or(false, |t| t.is_float());

                // TODO: We could handle `char` more gracefully.
                // TODO: Strings, though the lookup is a bit more hard (we need
                // to look at the canonical type of the pointee too, and check
                // is char, u8, or i8 I guess).
                let value = if is_integer {
                    let kind = match *canonical_ty.unwrap().kind() {
                        TypeKind::Int(kind) => kind,
                        _ => unreachable!(),
                    };

                    let mut val = cursor.evaluate().and_then(|v| v.as_int());
                    if val.is_none() || !kind.signedness_matches(val.unwrap()) {
                        let tu = ctx.translation_unit();
                        val = get_integer_literal_from_cursor(&cursor, tu);
                    }

                    val.map(|val| {
                        if kind == IntKind::Bool {
                            VarType::Bool(val != 0)
                        } else {
                            VarType::Int(val)
                        }
                    })
                } else if is_float {
                    cursor
                        .evaluate()
                        .and_then(|v| v.as_double())
                        .map(VarType::Float)
                } else {
                    cursor
                        .evaluate()
                        .and_then(|v| v.as_literal_string())
                        .map(VarType::String)
                };

                let mangling = cursor_mangling(ctx, &cursor);
                let var = Var::new(name, mangling, ty, value, is_const);

                Ok(ParseResult::New(var, Some(cursor)))
            }
            _ => {
                /* TODO */
                Err(ParseError::Continue)
            }
        }
    }
}

/// Try and parse a macro using all the macros parsed until now.
fn parse_macro(
    ctx: &BindgenContext,
    cursor: &clang::Cursor,
) -> Option<(Vec<u8>, cexpr::expr::EvalResult)> {
    use cexpr::expr;

    let cexpr_tokens = cursor.cexpr_tokens();

    let parser = expr::IdentifierParser::new(ctx.parsed_macros());

    match parser.macro_definition(&cexpr_tokens) {
        Ok((_, (id, val))) => Some((id.into(), val)),
        _ => None,
    }
}

fn parse_int_literal_tokens(cursor: &clang::Cursor) -> Option<i64> {
    use cexpr::expr;
    use cexpr::expr::EvalResult;

    let cexpr_tokens = cursor.cexpr_tokens();

    // TODO(emilio): We can try to parse other kinds of literals.
    match expr::expr(&cexpr_tokens) {
        Ok((_, EvalResult::Int(Wrapping(val)))) => Some(val),
        _ => None,
    }
}

fn get_integer_literal_from_cursor(
    cursor: &clang::Cursor,
    unit: &clang::TranslationUnit,
) -> Option<i64> {
    use clang_sys::*;
    let mut value = None;
    cursor.visit(|c| {
        match c.kind() {
            CXCursor_IntegerLiteral | CXCursor_UnaryOperator => {
                value = parse_int_literal_tokens(&c);
            }
            CXCursor_UnexposedExpr => {
                value = get_integer_literal_from_cursor(&c, unit);
            }
            _ => (),
        }
        if value.is_some() {
            CXChildVisit_Break
        } else {
            CXChildVisit_Continue
        }
    });
    value
}
<|endoftext|>


<commit_before>use attrs::AttributeSpanWrapper;
use field::Field;
use model::Model;
use proc_macro2::TokenStream;
use quote::quote_spanned;
use syn::{DeriveInput, Expr, Path, Type};
use util::{inner_of_option_ty, is_option_ty, wrap_in_dummy_mod};

pub fn derive(item: DeriveInput) -> TokenStream {
    let model = Model::from_item(&item, false);

    let treat_none_as_default_value = model.treat_none_as_default_value();
    let table_name = &model.table_name();
    let struct_name = &item.ident;

    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();

    let mut generate_borrowed_insert = true;

    let mut direct_field_ty = Vec::with_capacity(model.fields().len());
    let mut direct_field_assign = Vec::with_capacity(model.fields().len());
    let mut ref_field_ty = Vec::with_capacity(model.fields().len());
    let mut ref_field_assign = Vec::with_capacity(model.fields().len());

    for field in model.fields() {
        match (field.serialize_as.as_ref(), field.embed()) {
            (None, true) => {
                direct_field_ty.push(field_ty_embed(field, None));
                direct_field_assign.push(field_expr_embed(field, None));
                ref_field_ty.push(field_ty_embed(field, Some(quote!(&'insert))));
                ref_field_assign.push(field_expr_embed(field, Some(quote!(&))));
            }
            (None, false) => {
                direct_field_ty.push(field_ty(
                    field,
                    table_name,
                    None,
                    treat_none_as_default_value,
                ));
                direct_field_assign.push(field_expr(
                    field,
                    table_name,
                    None,
                    treat_none_as_default_value,
                ));
                ref_field_ty.push(field_ty(
                    field,
                    table_name,
                    Some(quote!(&'insert)),
                    treat_none_as_default_value,
                ));
                ref_field_assign.push(field_expr(
                    field,
                    table_name,
                    Some(quote!(&)),
                    treat_none_as_default_value,
                ));
            }
            (Some(AttributeSpanWrapper { item: ty, .. }), false) => {
                direct_field_ty.push(field_ty_serialize_as(
                    field,
                    table_name,
                    ty,
                    treat_none_as_default_value,
                ));
                direct_field_assign.push(field_expr_serialize_as(
                    field,
                    table_name,
                    ty,
                    treat_none_as_default_value,
                ));

                generate_borrowed_insert = false; // as soon as we hit one field with #[diesel(serialize_as)] there is no point in generating the impl of Insertable for borrowed structs
            }
            (Some(AttributeSpanWrapper { attribute_span, .. }), true) => {
                abort!(
                    attribute_span,
                    "`#[diesel(embed)]` cannot be combined with `#[diesel(serialize_as)]`"
                )
            }
        }
    }

    let insert_owned = quote! {
        impl #impl_generics Insertable<#table_name::table> for #struct_name #ty_generics
            #where_clause
        {
            type Values = <(#(#direct_field_ty,)*) as Insertable<#table_name::table>>::Values;

            fn values(self) -> <(#(#direct_field_ty,)*) as Insertable<#table_name::table>>::Values {
                (#(#direct_field_assign,)*).values()
            }
        }
    };

    let insert_borrowed = if generate_borrowed_insert {
        let mut impl_generics = item.generics.clone();
        impl_generics.params.push(parse_quote!('insert));
        let (impl_generics, ..) = impl_generics.split_for_impl();

        quote! {
            impl #impl_generics Insertable<#table_name::table>
                for &'insert #struct_name #ty_generics
            #where_clause
            {
                type Values = <(#(#ref_field_ty,)*) as Insertable<#table_name::table>>::Values;

                fn values(self) -> <(#(#ref_field_ty,)*) as Insertable<#table_name::table>>::Values {
                    (#(#ref_field_assign,)*).values()
                }
            }
        }
    } else {
        quote! {}
    };

    wrap_in_dummy_mod(quote! {
        use diesel::insertable::Insertable;
        use diesel::internal::derives::insertable::UndecoratedInsertRecord;
        use diesel::prelude::*;

        #[allow(unused_qualifications)]
        #insert_owned

	#[allow(unused_qualifications)]
        #insert_borrowed

        impl #impl_generics UndecoratedInsertRecord<#table_name::table>
                for #struct_name #ty_generics
            #where_clause
        {
        }
    })
}

fn field_ty_embed(field: &Field, lifetime: Option<TokenStream>) -> TokenStream {
    let field_ty = &field.ty;
    let span = field.span;
    quote_spanned!(span=> #lifetime #field_ty)
}

fn field_expr_embed(field: &Field, lifetime: Option<TokenStream>) -> TokenStream {
    let field_name = &field.name;
    quote!(#lifetime self.#field_name)
}

fn field_ty_serialize_as(
    field: &Field,
    table_name: &Path,
    ty: &Type,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let column_name = field.column_name();
    let span = field.span;
    if treat_none_as_default_value {
        let inner_ty = inner_of_option_ty(ty);

        quote_spanned! {span=>
            std::option::Option<diesel::dsl::Eq<
                #table_name::#column_name,
                #inner_ty,
            >>
        }
    } else {
        quote_spanned! {span=>
            diesel::dsl::Eq<
                #table_name::#column_name,
                #ty,
            >
        }
    }
}

fn field_expr_serialize_as(
    field: &Field,
    table_name: &Path,
    ty: &Type,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let field_name = &field.name;
    let column_name = field.column_name();
    let column = quote!(#table_name::#column_name);
    if treat_none_as_default_value {
        if is_option_ty(ty) {
            quote!(self.#field_name.map(|x| #column.eq(::std::convert::Into::<#ty>::into(x))))
        } else {
            quote!(std::option::Option::Some(#column.eq(::std::convert::Into::<#ty>::into(self.#field_name))))
        }
    } else {
        quote!(#column.eq(::std::convert::Into::<#ty>::into(self.#field_name)))
    }
}

fn field_ty(
    field: &Field,
    table_name: &Path,
    lifetime: Option<TokenStream>,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let column_name = field.column_name();
    let span = field.span;
    if treat_none_as_default_value {
        let inner_ty = inner_of_option_ty(&field.ty);

        quote_spanned! {span=>
            std::option::Option<diesel::dsl::Eq<
                #table_name::#column_name,
                #lifetime #inner_ty,
            >>
        }
    } else {
        let inner_ty = &field.ty;

        quote_spanned! {span=>
            diesel::dsl::Eq<
                #table_name::#column_name,
                #lifetime #inner_ty,
            >
        }
    }
}

fn field_expr(
    field: &Field,
    table_name: &Path,
    lifetime: Option<TokenStream>,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let field_name = &field.name;
    let column_name = field.column_name();

    let column: Expr = parse_quote!(#table_name::#column_name);
    if treat_none_as_default_value {
        if is_option_ty(&field.ty) {
            if lifetime.is_some() {
                quote!(self.#field_name.as_ref().map(|x| #column.eq(x)))
            } else {
                quote!(self.#field_name.map(|x| #column.eq(x)))
            }
        } else {
            quote!(std::option::Option::Some(#column.eq(#lifetime self.#field_name)))
        }
    } else {
        quote!(#column.eq(#lifetime self.#field_name))
    }
}
<commit_msg>fix fmt<commit_after>use attrs::AttributeSpanWrapper;
use field::Field;
use model::Model;
use proc_macro2::TokenStream;
use quote::quote_spanned;
use syn::{DeriveInput, Expr, Path, Type};
use util::{inner_of_option_ty, is_option_ty, wrap_in_dummy_mod};

pub fn derive(item: DeriveInput) -> TokenStream {
    let model = Model::from_item(&item, false);

    let treat_none_as_default_value = model.treat_none_as_default_value();
    let table_name = &model.table_name();
    let struct_name = &item.ident;

    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();

    let mut generate_borrowed_insert = true;

    let mut direct_field_ty = Vec::with_capacity(model.fields().len());
    let mut direct_field_assign = Vec::with_capacity(model.fields().len());
    let mut ref_field_ty = Vec::with_capacity(model.fields().len());
    let mut ref_field_assign = Vec::with_capacity(model.fields().len());

    for field in model.fields() {
        match (field.serialize_as.as_ref(), field.embed()) {
            (None, true) => {
                direct_field_ty.push(field_ty_embed(field, None));
                direct_field_assign.push(field_expr_embed(field, None));
                ref_field_ty.push(field_ty_embed(field, Some(quote!(&'insert))));
                ref_field_assign.push(field_expr_embed(field, Some(quote!(&))));
            }
            (None, false) => {
                direct_field_ty.push(field_ty(
                    field,
                    table_name,
                    None,
                    treat_none_as_default_value,
                ));
                direct_field_assign.push(field_expr(
                    field,
                    table_name,
                    None,
                    treat_none_as_default_value,
                ));
                ref_field_ty.push(field_ty(
                    field,
                    table_name,
                    Some(quote!(&'insert)),
                    treat_none_as_default_value,
                ));
                ref_field_assign.push(field_expr(
                    field,
                    table_name,
                    Some(quote!(&)),
                    treat_none_as_default_value,
                ));
            }
            (Some(AttributeSpanWrapper { item: ty, .. }), false) => {
                direct_field_ty.push(field_ty_serialize_as(
                    field,
                    table_name,
                    ty,
                    treat_none_as_default_value,
                ));
                direct_field_assign.push(field_expr_serialize_as(
                    field,
                    table_name,
                    ty,
                    treat_none_as_default_value,
                ));

                generate_borrowed_insert = false; // as soon as we hit one field with #[diesel(serialize_as)] there is no point in generating the impl of Insertable for borrowed structs
            }
            (Some(AttributeSpanWrapper { attribute_span, .. }), true) => {
                abort!(
                    attribute_span,
                    "`#[diesel(embed)]` cannot be combined with `#[diesel(serialize_as)]`"
                )
            }
        }
    }

    let insert_owned = quote! {
        impl #impl_generics Insertable<#table_name::table> for #struct_name #ty_generics
            #where_clause
        {
            type Values = <(#(#direct_field_ty,)*) as Insertable<#table_name::table>>::Values;

            fn values(self) -> <(#(#direct_field_ty,)*) as Insertable<#table_name::table>>::Values {
                (#(#direct_field_assign,)*).values()
            }
        }
    };

    let insert_borrowed = if generate_borrowed_insert {
        let mut impl_generics = item.generics.clone();
        impl_generics.params.push(parse_quote!('insert));
        let (impl_generics, ..) = impl_generics.split_for_impl();

        quote! {
            impl #impl_generics Insertable<#table_name::table>
                for &'insert #struct_name #ty_generics
            #where_clause
            {
                type Values = <(#(#ref_field_ty,)*) as Insertable<#table_name::table>>::Values;

                fn values(self) -> <(#(#ref_field_ty,)*) as Insertable<#table_name::table>>::Values {
                    (#(#ref_field_assign,)*).values()
                }
            }
        }
    } else {
        quote! {}
    };

    wrap_in_dummy_mod(quote! {
        use diesel::insertable::Insertable;
        use diesel::internal::derives::insertable::UndecoratedInsertRecord;
        use diesel::prelude::*;

        #[allow(unused_qualifications)]
        #insert_owned

    #[allow(unused_qualifications)]
        #insert_borrowed

        impl #impl_generics UndecoratedInsertRecord<#table_name::table>
                for #struct_name #ty_generics
            #where_clause
        {
        }
    })
}

fn field_ty_embed(field: &Field, lifetime: Option<TokenStream>) -> TokenStream {
    let field_ty = &field.ty;
    let span = field.span;
    quote_spanned!(span=> #lifetime #field_ty)
}

fn field_expr_embed(field: &Field, lifetime: Option<TokenStream>) -> TokenStream {
    let field_name = &field.name;
    quote!(#lifetime self.#field_name)
}

fn field_ty_serialize_as(
    field: &Field,
    table_name: &Path,
    ty: &Type,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let column_name = field.column_name();
    let span = field.span;
    if treat_none_as_default_value {
        let inner_ty = inner_of_option_ty(ty);

        quote_spanned! {span=>
            std::option::Option<diesel::dsl::Eq<
                #table_name::#column_name,
                #inner_ty,
            >>
        }
    } else {
        quote_spanned! {span=>
            diesel::dsl::Eq<
                #table_name::#column_name,
                #ty,
            >
        }
    }
}

fn field_expr_serialize_as(
    field: &Field,
    table_name: &Path,
    ty: &Type,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let field_name = &field.name;
    let column_name = field.column_name();
    let column = quote!(#table_name::#column_name);
    if treat_none_as_default_value {
        if is_option_ty(ty) {
            quote!(self.#field_name.map(|x| #column.eq(::std::convert::Into::<#ty>::into(x))))
        } else {
            quote!(std::option::Option::Some(#column.eq(::std::convert::Into::<#ty>::into(self.#field_name))))
        }
    } else {
        quote!(#column.eq(::std::convert::Into::<#ty>::into(self.#field_name)))
    }
}

fn field_ty(
    field: &Field,
    table_name: &Path,
    lifetime: Option<TokenStream>,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let column_name = field.column_name();
    let span = field.span;
    if treat_none_as_default_value {
        let inner_ty = inner_of_option_ty(&field.ty);

        quote_spanned! {span=>
            std::option::Option<diesel::dsl::Eq<
                #table_name::#column_name,
                #lifetime #inner_ty,
            >>
        }
    } else {
        let inner_ty = &field.ty;

        quote_spanned! {span=>
            diesel::dsl::Eq<
                #table_name::#column_name,
                #lifetime #inner_ty,
            >
        }
    }
}

fn field_expr(
    field: &Field,
    table_name: &Path,
    lifetime: Option<TokenStream>,
    treat_none_as_default_value: bool,
) -> TokenStream {
    let field_name = &field.name;
    let column_name = field.column_name();

    let column: Expr = parse_quote!(#table_name::#column_name);
    if treat_none_as_default_value {
        if is_option_ty(&field.ty) {
            if lifetime.is_some() {
                quote!(self.#field_name.as_ref().map(|x| #column.eq(x)))
            } else {
                quote!(self.#field_name.map(|x| #column.eq(x)))
            }
        } else {
            quote!(std::option::Option::Some(#column.eq(#lifetime self.#field_name)))
        }
    } else {
        quote!(#column.eq(#lifetime self.#field_name))
    }
}
<|endoftext|>


<commit_before>// Copyright 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
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::TokenStream;
use quote::quote;
use syn::{
    parse::Parser,
    parse_quote,
    punctuated::Punctuated,
    token::{Comma, Unsafe},
    Attribute, FnArg, ForeignItem, Ident, ImplItem, Item, ReturnType,
};

use super::{
    unqualify::{unqualify_params, unqualify_ret_type},
    RsCodegenResult, Use,
};
use crate::{
    conversion::{
        analysis::fun::{
            ArgumentAnalysis, FnAnalysis, FnKind, MethodKind, RustRenameStrategy,
            TraitMethodDetails, UnsafetyNeeded,
        },
        api::ImplBlockDetails,
        codegen_rs::lifetime::add_lifetime_to_all_params,
    },
    types::{Namespace, QualifiedName},
};
use crate::{
    conversion::{api::FuncToConvert, codegen_rs::lifetime::add_explicit_lifetime_if_necessary},
    types::make_ident,
};

impl UnsafetyNeeded {
    fn bridge_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::None => None,
            _ => Some(parse_quote! { unsafe }),
        }
    }

    fn wrapper_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::Always => Some(parse_quote! { unsafe }),
            _ => None,
        }
    }
}

pub(super) fn gen_function(
    ns: &Namespace,
    fun: FuncToConvert,
    analysis: FnAnalysis,
    cpp_call_name: String,
) -> RsCodegenResult {
    if analysis.ignore_reason.is_err() || !analysis.externally_callable {
        return RsCodegenResult::default();
    }
    let cxxbridge_name = analysis.cxxbridge_name;
    let rust_name = analysis.rust_name;
    let ret_type = analysis.ret_type;
    let param_details = analysis.param_details;
    let wrapper_function_needed = analysis.cpp_wrapper.is_some();
    let params = analysis.params;
    let vis = analysis.vis;
    let kind = analysis.kind;
    let doc_attr = fun.doc_attr;

    let mut cpp_name_attr = Vec::new();
    let mut impl_entry = None;
    let rust_name_attr: Vec<_> = match &analysis.rust_rename_strategy {
        RustRenameStrategy::RenameUsingRustAttr => Attribute::parse_outer
            .parse2(quote!(
                #[rust_name = #rust_name]
            ))
            .unwrap(),
        _ => Vec::new(),
    };
    let wrapper_unsafety = analysis.requires_unsafe.wrapper_token();
    let fn_generator = FnGenerator {
        param_details: &param_details,
        cxxbridge_name: &cxxbridge_name,
        rust_name: &rust_name,
        unsafety: &wrapper_unsafety,
        doc_attr: &doc_attr,
    };
    let (mut materialization, is_a_trait) = match kind {
        FnKind::Method(..) => (None, false),
        FnKind::Function => match analysis.rust_rename_strategy {
            RustRenameStrategy::RenameInOutputMod(alias) => {
                (Some(Use::UsedFromCxxBridgeWithAlias(alias)), false)
            }
            _ => (Some(Use::UsedFromCxxBridge), false),
        },
        FnKind::TraitMethod { ref details, .. } => {
            (fn_generator.generate_trait_impl(details, &ret_type), true)
        }
    };
    let any_param_needs_rust_conversion = param_details
        .iter()
        .any(|pd| pd.conversion.rust_work_needed());
    let rust_wrapper_needed = any_param_needs_rust_conversion
        || (cxxbridge_name != rust_name && matches!(kind, FnKind::Method(..)));
    if rust_wrapper_needed && !is_a_trait {
        match kind {
            FnKind::Method(ref type_name, MethodKind::Constructor) => {
                // Constructor.
                impl_entry = Some(fn_generator.generate_constructor_impl(type_name));
            }
            FnKind::Method(ref type_name, ref method_kind) => {
                // Method, or static method.
                impl_entry = Some(fn_generator.generate_method_impl(
                    matches!(
                        method_kind,
                        MethodKind::MakeUnique | MethodKind::Constructor
                    ),
                    type_name,
                    &ret_type,
                ));
            }
            _ => {
                // Generate plain old function
                materialization = Some(Use::Custom(fn_generator.generate_function_impl(&ret_type)));
            }
        }
    }
    if cxxbridge_name != cpp_call_name && !wrapper_function_needed {
        cpp_name_attr = Attribute::parse_outer
            .parse2(quote!(
                #[cxx_name = #cpp_call_name]
            ))
            .unwrap();
    }
    // In very rare occasions, we might need to give an explicit lifetime.
    let (lifetime_tokens, params, ret_type) =
        add_explicit_lifetime_if_necessary(&param_details, params, &ret_type);

    // Finally - namespace support. All the Types in everything
    // above this point are fully qualified. We need to unqualify them.
    // We need to do that _after_ the above wrapper_function_needed
    // work, because it relies upon spotting fully qualified names like
    // std::unique_ptr. However, after it's done its job, all such
    // well-known types should be unqualified already (e.g. just UniquePtr)
    // and the following code will act to unqualify only those types
    // which the user has declared.
    let params = unqualify_params(params);
    let ret_type = unqualify_ret_type(ret_type.into_owned());
    // And we need to make an attribute for the namespace that the function
    // itself is in.
    let namespace_attr = if ns.is_empty() || wrapper_function_needed {
        Vec::new()
    } else {
        let namespace_string = ns.to_string();
        Attribute::parse_outer
            .parse2(quote!(
                #[namespace = #namespace_string]
            ))
            .unwrap()
    };
    // At last, actually generate the cxx::bridge entry.
    let bridge_unsafety = analysis.requires_unsafe.bridge_token();
    let extern_c_mod_item = ForeignItem::Fn(parse_quote!(
        #(#namespace_attr)*
        #(#rust_name_attr)*
        #(#cpp_name_attr)*
        #doc_attr
        #vis #bridge_unsafety fn #cxxbridge_name #lifetime_tokens ( #params ) #ret_type;
    ));
    RsCodegenResult {
        extern_c_mod_items: vec![extern_c_mod_item],
        bridge_items: Vec::new(),
        global_items: Vec::new(),
        bindgen_mod_items: Vec::new(),
        impl_entry,
        materializations: materialization.into_iter().collect(),
        extern_rust_mod_items: Vec::new(),
    }
}

/// Knows how to generate a given function.
#[derive(Clone)]
struct FnGenerator<'a> {
    param_details: &'a [ArgumentAnalysis],
    cxxbridge_name: &'a Ident,
    rust_name: &'a str,
    unsafety: &'a Option<Unsafe>,
    doc_attr: &'a Option<Attribute>,
}

impl<'a> FnGenerator<'a> {
    fn generate_arg_lists(&self, avoid_self: bool) -> (Punctuated<FnArg, Comma>, Vec<TokenStream>) {
        let mut wrapper_params: Punctuated<FnArg, Comma> = Punctuated::new();
        let mut arg_list = Vec::new();

        for pd in self.param_details {
            let type_name = pd.conversion.rust_wrapper_unconverted_type();
            let wrapper_arg_name = if pd.self_type.is_some() && !avoid_self {
                parse_quote!(self)
            } else {
                pd.name.clone()
            };
            let param_mutability = pd.conversion.rust_conversion.requires_mutability();
            wrapper_params.push(parse_quote!(
                #param_mutability #wrapper_arg_name: #type_name
            ));
            arg_list.push(pd.conversion.rust_conversion(wrapper_arg_name));
        }
        (wrapper_params, arg_list)
    }

    /// Generate an 'impl Type { methods-go-here }' item
    fn generate_method_impl(
        &self,
        avoid_self: bool,
        impl_block_type_name: &QualifiedName,
        ret_type: &ReturnType,
    ) -> Box<ImplBlockDetails> {
        let (wrapper_params, arg_list) = self.generate_arg_lists(avoid_self);
        let (lifetime_tokens, wrapper_params, ret_type) =
            add_explicit_lifetime_if_necessary(self.param_details, wrapper_params, ret_type);
        let rust_name = make_ident(self.rust_name);
        let unsafety = self.unsafety;
        let doc_attr = self.doc_attr;
        let cxxbridge_name = self.cxxbridge_name;
        Box::new(ImplBlockDetails {
            item: ImplItem::Method(parse_quote! {
                #doc_attr
                pub #unsafety fn #rust_name #lifetime_tokens ( #wrapper_params ) #ret_type {
                    cxxbridge::#cxxbridge_name ( #(#arg_list),* )
                }
            }),
            ty: impl_block_type_name.get_final_ident(),
        })
    }

    /// Generate an 'impl Trait for Type { methods-go-here }' in its entrety.
    fn generate_trait_impl(
        &self,
        details: &TraitMethodDetails,
        ret_type: &ReturnType,
    ) -> Option<Use> {
        let (mut wrapper_params, arg_list) = self.generate_arg_lists(details.avoid_self);
        if let Some(parameter_reordering) = &details.parameter_reordering {
            wrapper_params = Self::reorder_parameters(wrapper_params, parameter_reordering);
        }
        let (lifetime_tokens, wrapper_params, ret_type) =
            add_explicit_lifetime_if_necessary(self.param_details, wrapper_params, ret_type);
        let doc_attr = self.doc_attr;
        let unsafety = self.unsafety;
        let cxxbridge_name = self.cxxbridge_name;
        let trait_signature = &details.trait_signature;
        let trait_unsafety = &details.trait_unsafety;
        let impl_for_specifics = &details.impl_for_specifics;
        let method_name = &details.method_name;
        let call_body = quote! {
            cxxbridge::#cxxbridge_name ( #(#arg_list),* )
        };
        let call_body = if details.trait_call_is_unsafe {
            quote! {
                unsafe {
                    #call_body
                }
            }
        } else {
            call_body
        };
        Some(Use::Custom(Box::new(parse_quote! {
            #trait_unsafety impl #lifetime_tokens #trait_signature for #impl_for_specifics {
                #doc_attr
                #unsafety fn #method_name ( #wrapper_params ) #ret_type {
                    #call_body
                }
            }
        })))
    }

    /// Generate a 'impl Type { methods-go-here }' item which is a constructor
    /// for use with moveit traits.
    fn generate_constructor_impl(
        &self,
        impl_block_type_name: &QualifiedName,
    ) -> Box<ImplBlockDetails> {
        let (wrapper_params, arg_list) = self.generate_arg_lists(true);
        let mut wrapper_params: Punctuated<FnArg, Comma> =
            wrapper_params.into_iter().skip(1).collect();
        let ptr_arg_name = &arg_list[0];
        let rust_name = make_ident(&self.rust_name);
        let any_references = self.param_details.iter().any(|pd| pd.was_reference);
        let (lifetime_param, lifetime_addition) = if any_references {
            add_lifetime_to_all_params(&mut wrapper_params);
            (quote! { <'a> }, quote! { + 'a })
        } else {
            (quote! {}, quote! {})
        };
        let cxxbridge_name = self.cxxbridge_name;
        let body = quote! {
            autocxx::moveit::new::by_raw(move |#ptr_arg_name| {
                let #ptr_arg_name = #ptr_arg_name.get_unchecked_mut().as_mut_ptr();
                cxxbridge::#cxxbridge_name(#(#arg_list),* )
            })
        };
        let body = if self.unsafety.is_some() {
            // No need for `unsafe` inside the function
            body
        } else {
            quote! {
                unsafe { #body }
            }
        };
        let doc_attr = self.doc_attr;
        let unsafety = self.unsafety;
        Box::new(ImplBlockDetails {
            item: ImplItem::Method(parse_quote! {
                #doc_attr
                pub #unsafety fn #rust_name #lifetime_param ( #wrapper_params ) -> impl autocxx::moveit::new::New<Output=Self> #lifetime_addition {
                    #body
                }
            }),
            ty: impl_block_type_name.get_final_ident(),
        })
    }

    /// Generate a function call wrapper
    fn generate_function_impl(&self, ret_type: &ReturnType) -> Box<Item> {
        let (wrapper_params, arg_list) = self.generate_arg_lists(false);
        let rust_name = make_ident(self.rust_name);
        let doc_attr = self.doc_attr;
        let unsafety = self.unsafety;
        Box::new(Item::Fn(parse_quote! {
            #doc_attr
            pub #unsafety fn #rust_name ( #wrapper_params ) #ret_type {
                cxxbridge::#rust_name ( #(#arg_list),* )
            }
        }))
    }

    fn reorder_parameters(
        params: Punctuated<FnArg, Comma>,
        parameter_ordering: &[usize],
    ) -> Punctuated<FnArg, Comma> {
        let old_params = params.into_iter().collect::<Vec<_>>();
        parameter_ordering
            .iter()
            .map(|n| old_params.get(*n).unwrap().clone())
            .collect()
    }
}
<commit_msg>Locating trait items in root.<commit_after>// Copyright 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
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::TokenStream;
use quote::quote;
use syn::{
    parse::Parser,
    parse_quote,
    punctuated::Punctuated,
    token::{Comma, Unsafe},
    Attribute, FnArg, ForeignItem, Ident, ImplItem, Item, ReturnType,
};

use super::{
    unqualify::{unqualify_params, unqualify_ret_type},
    RsCodegenResult, Use,
};
use crate::{
    conversion::{
        analysis::fun::{
            ArgumentAnalysis, FnAnalysis, FnKind, MethodKind, RustRenameStrategy,
            TraitMethodDetails, UnsafetyNeeded,
        },
        api::ImplBlockDetails,
        codegen_rs::lifetime::add_lifetime_to_all_params,
    },
    types::{Namespace, QualifiedName},
};
use crate::{
    conversion::{api::FuncToConvert, codegen_rs::lifetime::add_explicit_lifetime_if_necessary},
    types::make_ident,
};

impl UnsafetyNeeded {
    fn bridge_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::None => None,
            _ => Some(parse_quote! { unsafe }),
        }
    }

    fn wrapper_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::Always => Some(parse_quote! { unsafe }),
            _ => None,
        }
    }
}

pub(super) fn gen_function(
    ns: &Namespace,
    fun: FuncToConvert,
    analysis: FnAnalysis,
    cpp_call_name: String,
) -> RsCodegenResult {
    if analysis.ignore_reason.is_err() || !analysis.externally_callable {
        return RsCodegenResult::default();
    }
    let cxxbridge_name = analysis.cxxbridge_name;
    let rust_name = analysis.rust_name;
    let ret_type = analysis.ret_type;
    let param_details = analysis.param_details;
    let wrapper_function_needed = analysis.cpp_wrapper.is_some();
    let params = analysis.params;
    let vis = analysis.vis;
    let kind = analysis.kind;
    let doc_attr = fun.doc_attr;

    let mut cpp_name_attr = Vec::new();
    let mut impl_entry = None;
    let mut trait_impl_entry = None;
    let rust_name_attr: Vec<_> = match &analysis.rust_rename_strategy {
        RustRenameStrategy::RenameUsingRustAttr => Attribute::parse_outer
            .parse2(quote!(
                #[rust_name = #rust_name]
            ))
            .unwrap(),
        _ => Vec::new(),
    };
    let wrapper_unsafety = analysis.requires_unsafe.wrapper_token();
    let fn_generator = FnGenerator {
        param_details: &param_details,
        cxxbridge_name: &cxxbridge_name,
        rust_name: &rust_name,
        unsafety: &wrapper_unsafety,
        doc_attr: &doc_attr,
    };
    let mut materialization = match kind {
        FnKind::Method(..) | FnKind::TraitMethod { .. } => None,
        FnKind::Function => match analysis.rust_rename_strategy {
            RustRenameStrategy::RenameInOutputMod(alias) => {
                Some(Use::UsedFromCxxBridgeWithAlias(alias))
            }
            _ => Some(Use::UsedFromCxxBridge),
        },
    };
    let any_param_needs_rust_conversion = param_details
        .iter()
        .any(|pd| pd.conversion.rust_work_needed());
    let rust_wrapper_needed = any_param_needs_rust_conversion
        || (cxxbridge_name != rust_name && matches!(kind, FnKind::Method(..)));
    if rust_wrapper_needed {
        match kind {
            FnKind::Method(ref type_name, MethodKind::Constructor) => {
                // Constructor.
                impl_entry = Some(fn_generator.generate_constructor_impl(type_name));
            }
            FnKind::Method(ref type_name, ref method_kind) => {
                // Method, or static method.
                impl_entry = Some(fn_generator.generate_method_impl(
                    matches!(
                        method_kind,
                        MethodKind::MakeUnique | MethodKind::Constructor
                    ),
                    type_name,
                    &ret_type,
                ));
            }
            FnKind::TraitMethod { ref details, .. } => {
                trait_impl_entry = Some(fn_generator.generate_trait_impl(details, &ret_type));
            }
            _ => {
                // Generate plain old function
                materialization = Some(Use::Custom(fn_generator.generate_function_impl(&ret_type)));
            }
        }
    }
    if cxxbridge_name != cpp_call_name && !wrapper_function_needed {
        cpp_name_attr = Attribute::parse_outer
            .parse2(quote!(
                #[cxx_name = #cpp_call_name]
            ))
            .unwrap();
    }
    // In very rare occasions, we might need to give an explicit lifetime.
    let (lifetime_tokens, params, ret_type) =
        add_explicit_lifetime_if_necessary(&param_details, params, &ret_type);

    // Finally - namespace support. All the Types in everything
    // above this point are fully qualified. We need to unqualify them.
    // We need to do that _after_ the above wrapper_function_needed
    // work, because it relies upon spotting fully qualified names like
    // std::unique_ptr. However, after it's done its job, all such
    // well-known types should be unqualified already (e.g. just UniquePtr)
    // and the following code will act to unqualify only those types
    // which the user has declared.
    let params = unqualify_params(params);
    let ret_type = unqualify_ret_type(ret_type.into_owned());
    // And we need to make an attribute for the namespace that the function
    // itself is in.
    let namespace_attr = if ns.is_empty() || wrapper_function_needed {
        Vec::new()
    } else {
        let namespace_string = ns.to_string();
        Attribute::parse_outer
            .parse2(quote!(
                #[namespace = #namespace_string]
            ))
            .unwrap()
    };
    // At last, actually generate the cxx::bridge entry.
    let bridge_unsafety = analysis.requires_unsafe.bridge_token();
    let extern_c_mod_item = ForeignItem::Fn(parse_quote!(
        #(#namespace_attr)*
        #(#rust_name_attr)*
        #(#cpp_name_attr)*
        #doc_attr
        #vis #bridge_unsafety fn #cxxbridge_name #lifetime_tokens ( #params ) #ret_type;
    ));
    RsCodegenResult {
        extern_c_mod_items: vec![extern_c_mod_item],
        bridge_items: Vec::new(),
        global_items: trait_impl_entry.into_iter().collect(),
        bindgen_mod_items: Vec::new(),
        impl_entry,
        materializations: materialization.into_iter().collect(),
        extern_rust_mod_items: Vec::new(),
    }
}

/// Knows how to generate a given function.
#[derive(Clone)]
struct FnGenerator<'a> {
    param_details: &'a [ArgumentAnalysis],
    cxxbridge_name: &'a Ident,
    rust_name: &'a str,
    unsafety: &'a Option<Unsafe>,
    doc_attr: &'a Option<Attribute>,
}

impl<'a> FnGenerator<'a> {
    fn generate_arg_lists(&self, avoid_self: bool) -> (Punctuated<FnArg, Comma>, Vec<TokenStream>) {
        let mut wrapper_params: Punctuated<FnArg, Comma> = Punctuated::new();
        let mut arg_list = Vec::new();

        for pd in self.param_details {
            let type_name = pd.conversion.rust_wrapper_unconverted_type();
            let wrapper_arg_name = if pd.self_type.is_some() && !avoid_self {
                parse_quote!(self)
            } else {
                pd.name.clone()
            };
            let param_mutability = pd.conversion.rust_conversion.requires_mutability();
            wrapper_params.push(parse_quote!(
                #param_mutability #wrapper_arg_name: #type_name
            ));
            arg_list.push(pd.conversion.rust_conversion(wrapper_arg_name));
        }
        (wrapper_params, arg_list)
    }

    /// Generate an 'impl Type { methods-go-here }' item
    fn generate_method_impl(
        &self,
        avoid_self: bool,
        impl_block_type_name: &QualifiedName,
        ret_type: &ReturnType,
    ) -> Box<ImplBlockDetails> {
        let (wrapper_params, arg_list) = self.generate_arg_lists(avoid_self);
        let (lifetime_tokens, wrapper_params, ret_type) =
            add_explicit_lifetime_if_necessary(self.param_details, wrapper_params, ret_type);
        let rust_name = make_ident(self.rust_name);
        let unsafety = self.unsafety;
        let doc_attr = self.doc_attr;
        let cxxbridge_name = self.cxxbridge_name;
        Box::new(ImplBlockDetails {
            item: ImplItem::Method(parse_quote! {
                #doc_attr
                pub #unsafety fn #rust_name #lifetime_tokens ( #wrapper_params ) #ret_type {
                    cxxbridge::#cxxbridge_name ( #(#arg_list),* )
                }
            }),
            ty: impl_block_type_name.get_final_ident(),
        })
    }

    /// Generate an 'impl Trait for Type { methods-go-here }' in its entrety.
    fn generate_trait_impl(&self, details: &TraitMethodDetails, ret_type: &ReturnType) -> Item {
        let (mut wrapper_params, arg_list) = self.generate_arg_lists(details.avoid_self);
        if let Some(parameter_reordering) = &details.parameter_reordering {
            wrapper_params = Self::reorder_parameters(wrapper_params, parameter_reordering);
        }
        let (lifetime_tokens, wrapper_params, ret_type) =
            add_explicit_lifetime_if_necessary(self.param_details, wrapper_params, ret_type);
        let doc_attr = self.doc_attr;
        let unsafety = self.unsafety;
        let cxxbridge_name = self.cxxbridge_name;
        let trait_signature = &details.trait_signature;
        let trait_unsafety = &details.trait_unsafety;
        let impl_for_specifics = &details.impl_for_specifics;
        let method_name = &details.method_name;
        let call_body = quote! {
            cxxbridge::#cxxbridge_name ( #(#arg_list),* )
        };
        let call_body = if details.trait_call_is_unsafe {
            quote! {
                unsafe {
                    #call_body
                }
            }
        } else {
            call_body
        };
        parse_quote! {
            #trait_unsafety impl #lifetime_tokens #trait_signature for #impl_for_specifics {
                #doc_attr
                #unsafety fn #method_name ( #wrapper_params ) #ret_type {
                    #call_body
                }
            }
        }
    }

    /// Generate a 'impl Type { methods-go-here }' item which is a constructor
    /// for use with moveit traits.
    fn generate_constructor_impl(
        &self,
        impl_block_type_name: &QualifiedName,
    ) -> Box<ImplBlockDetails> {
        let (wrapper_params, arg_list) = self.generate_arg_lists(true);
        let mut wrapper_params: Punctuated<FnArg, Comma> =
            wrapper_params.into_iter().skip(1).collect();
        let ptr_arg_name = &arg_list[0];
        let rust_name = make_ident(&self.rust_name);
        let any_references = self.param_details.iter().any(|pd| pd.was_reference);
        let (lifetime_param, lifetime_addition) = if any_references {
            add_lifetime_to_all_params(&mut wrapper_params);
            (quote! { <'a> }, quote! { + 'a })
        } else {
            (quote! {}, quote! {})
        };
        let cxxbridge_name = self.cxxbridge_name;
        let body = quote! {
            autocxx::moveit::new::by_raw(move |#ptr_arg_name| {
                let #ptr_arg_name = #ptr_arg_name.get_unchecked_mut().as_mut_ptr();
                cxxbridge::#cxxbridge_name(#(#arg_list),* )
            })
        };
        let body = if self.unsafety.is_some() {
            // No need for `unsafe` inside the function
            body
        } else {
            quote! {
                unsafe { #body }
            }
        };
        let doc_attr = self.doc_attr;
        let unsafety = self.unsafety;
        Box::new(ImplBlockDetails {
            item: ImplItem::Method(parse_quote! {
                #doc_attr
                pub #unsafety fn #rust_name #lifetime_param ( #wrapper_params ) -> impl autocxx::moveit::new::New<Output=Self> #lifetime_addition {
                    #body
                }
            }),
            ty: impl_block_type_name.get_final_ident(),
        })
    }

    /// Generate a function call wrapper
    fn generate_function_impl(&self, ret_type: &ReturnType) -> Box<Item> {
        let (wrapper_params, arg_list) = self.generate_arg_lists(false);
        let rust_name = make_ident(self.rust_name);
        let doc_attr = self.doc_attr;
        let unsafety = self.unsafety;
        Box::new(Item::Fn(parse_quote! {
            #doc_attr
            pub #unsafety fn #rust_name ( #wrapper_params ) #ret_type {
                cxxbridge::#rust_name ( #(#arg_list),* )
            }
        }))
    }

    fn reorder_parameters(
        params: Punctuated<FnArg, Comma>,
        parameter_ordering: &[usize],
    ) -> Punctuated<FnArg, Comma> {
        let old_params = params.into_iter().collect::<Vec<_>>();
        parameter_ordering
            .iter()
            .map(|n| old_params.get(*n).unwrap().clone())
            .collect()
    }
}
<|endoftext|>