<commit_before>use super::{Tokenizer, Whitespace, CompareStrs};
use super::{ScanError, OtherScanError};

use std::fmt::{Show, Formatter, FormatError};
use std::str::CharRange;

pub trait ScanCursor<'scanee>: Clone + Eq {
	fn expect_tok(&self, s: &str) -> Result<Self, ScanError>;
	fn consumed(&self) -> uint;
	fn pop_token(&self) -> Option<(&'scanee str, Self)>;
	fn pop_ws(&self) -> Self;
	fn slice_from(&self, from: uint) -> Self;
	fn str_slice_to(&self, to: uint) -> &'scanee str;
	fn str_slice_to_cur(&self, to: &Self) -> &'scanee str;
	fn tail_str(&self) -> &'scanee str;
	fn is_empty(&self) -> bool;
	fn compare_strs(&self, a: &str, b: &str) -> bool;

	fn expect_eof(&self) -> Result<(), ScanError> {
		if self.pop_token().is_some() {
			Err(self.expected_eof())
		} else {
			Ok(())
		}
	}

	fn expected(&self, desc: &str) -> ScanError {
		let msg = match self.pop_token() {
			Some((got, _)) => format!("expected {}, got `{}`", desc, got),
			None => format!("expected {}, got end of input", desc)
		};

		OtherScanError(msg, self.consumed())
	}

	fn expected_tok(&self, tok: &str) -> ScanError {
		let toks = [tok];
		self.expected_one_of(&toks)
	}

	fn expected_eof(&self) -> ScanError {
		self.expected_one_of(&[])
	}

	fn expected_one_of(&self, toks: &[&str]) -> ScanError {
		let mut toks = toks.iter().map(|s| format!("`{}`", s));
		let toks = {
			let first = toks.next();
			first.map(|first| toks.fold(first, |a,b| format!("{}, {}", a, b)))
		};

		let msg = match (toks, self.pop_token()) {
			(Some(exp), Some((got, _))) => format!("expected {}, got `{}`", exp, got),
			(Some(exp), None) => format!("expected {}, got end of input", exp),
			(None, Some((got, _))) => format!("expected end of input, got `{}`", got),
			(None, None) => "expected end of input".into_string()
		};

		OtherScanError(msg, self.consumed())
	}

	fn expected_min_repeats(&self, min: uint, got: uint) -> ScanError {
		OtherScanError(format!("expected at least {} repeats, got {}", min, got), self.consumed())
	}
}

#[deriving(Clone, Eq, PartialEq)]
pub struct Cursor<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> {
	slice: &'a str,
	offset: uint,
	tc: Tok,
	sp: Sp,
	cs: Cs,
}

impl<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> Cursor<'a, Tok, Sp, Cs> {
	pub fn new<'b>(s: &'b str, tc: Tok, sp: Sp, cs: Cs) -> Cursor<'b, Tok, Sp, Cs> {
		Cursor {
			slice: s,
			offset: 0,
			tc: tc,
			sp: sp,
			cs: cs,
		}
	}
}

impl<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> Show for Cursor<'a, Tok, Sp, Cs> {
	fn fmt(&self, f: &mut Formatter) -> Result<(), FormatError> {
		try!(write!(f, "Cursor<{}, {}, {}> {{ offset: {}, .. }}", self.tc, self.sp, self.cs, self.offset));
		Ok(())
	}
}

impl<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> ScanCursor<'a> for Cursor<'a, Tok, Sp, Cs> {
	fn expect_tok(&self, s: &str) -> Result<Cursor<'a, Tok, Sp, Cs>, ScanError> {
		debug!("{}.expect_tok({})", self, s);
		match self.pop_token() {
			Some((tok, ref cur)) if self.compare_strs(s, tok) => Ok(cur.clone()),
			_ => Err(self.expected_tok(s))
		}
	}

	fn consumed(&self) -> uint {
		self.offset
	}

	fn pop_token(&self) -> Option<(&'a str, Cursor<'a, Tok, Sp, Cs>)> {
		debug!("{}.pop_token()", self);
		// First, strip out leading whitespace.  It's up to the whitespace policy to *not* strip characters it wants to turn into a token.
		let cur = self.pop_ws();

		// Next, check to see if there is a whitespace token.  This allows the space policy to do things like ignore most whitespace, but turn line breaks into explicit tokens.  Note that unlike the regular Tokenizer, the Whitespace policy is responsible for returning the str slice itself.  This is used to do things like map all whitespace to a single `" "` token.
		match self.sp.token_len(cur.tail_str()) {
			Some(end) => return Some((cur.str_slice_to(end), cur.slice_from(end))),
			None => ()
		}

		// Do not assume that empty input means we can't match a token; the token class might, for example, turn end-of-input into an explicit token.
		let tail_str = cur.tail_str();
		match self.tc.token_len(tail_str) {
			Some(end) => Some((cur.str_slice_to(end), cur.slice_from(end))),
			None => {
				// One of two things: either we have some input left and will thus return a single-character token, or there is nothing left whereby we return None.
				if cur.is_empty() {
					return None;
				} else {
					let CharRange { ch: _, next } = tail_str.char_range_at(0);
					Some((cur.str_slice_to(next), cur.slice_from(next)))
				}
			},
		}
	}

	fn pop_ws(&self) -> Cursor<'a, Tok, Sp, Cs> {
		debug!("{}.pop_ws()", self);

		self.slice_from(self.sp.strip_len(self.tail_str()))
	}

	fn slice_from(&self, from: uint) -> Cursor<'a, Tok, Sp, Cs> {
		Cursor {
			offset: ::std::cmp::min(self.slice.len(), self.offset + from),
			..self.clone()
		}
	}

	fn str_slice_to(&self, to: uint) -> &'a str {
		self.tail_str().slice_to(to)
	}

	fn str_slice_to_cur(&self, to: &Cursor<'a, Tok, Sp, Cs>) -> &'a str {
		self.slice.slice(self.offset, to.offset)
	}

	fn tail_str(&self) -> &'a str {
		self.slice.slice_from(self.offset)
	}

	fn is_empty(&self) -> bool {
		self.offset == self.slice.len()
	}

	fn compare_strs(&self, a: &str, b: &str) -> bool {
		self.cs.compare_strs(a, b)
	}
}
<commit_msg>Cursor now escapes strings in output.<commit_after>use super::{Tokenizer, Whitespace, CompareStrs};
use super::{ScanError, OtherScanError};

use std::fmt::{Show, Formatter, FormatError};
use std::str::CharRange;

pub trait ScanCursor<'scanee>: Clone + Eq {
	fn expect_tok(&self, s: &str) -> Result<Self, ScanError>;
	fn consumed(&self) -> uint;
	fn pop_token(&self) -> Option<(&'scanee str, Self)>;
	fn pop_ws(&self) -> Self;
	fn slice_from(&self, from: uint) -> Self;
	fn str_slice_to(&self, to: uint) -> &'scanee str;
	fn str_slice_to_cur(&self, to: &Self) -> &'scanee str;
	fn tail_str(&self) -> &'scanee str;
	fn is_empty(&self) -> bool;
	fn compare_strs(&self, a: &str, b: &str) -> bool;

	fn expect_eof(&self) -> Result<(), ScanError> {
		if self.pop_token().is_some() {
			Err(self.expected_eof())
		} else {
			Ok(())
		}
	}

	fn expected(&self, desc: &str) -> ScanError {
		let msg = match self.pop_token() {
			Some((got, _)) => format!("expected {}, got `{}`", desc, got.escape_default()),
			None => format!("expected {}, got end of input", desc)
		};

		OtherScanError(msg, self.consumed())
	}

	fn expected_tok(&self, tok: &str) -> ScanError {
		let toks = [tok];
		self.expected_one_of(&toks)
	}

	fn expected_eof(&self) -> ScanError {
		self.expected_one_of(&[])
	}

	fn expected_one_of(&self, toks: &[&str]) -> ScanError {
		let mut toks = toks.iter().map(|s| format!("`{}`", s.escape_default()));
		let toks = {
			let first = toks.next();
			first.map(|first| toks.fold(first, |a,b| format!("{}, {}", a, b)))
		};

		let msg = match (toks, self.pop_token()) {
			(Some(exp), Some((got, _))) => format!("expected {}, got `{}`", exp, got.escape_default()),
			(Some(exp), None) => format!("expected {}, got end of input", exp),
			(None, Some((got, _))) => format!("expected end of input, got `{}`", got.escape_default()),
			(None, None) => "expected end of input".into_string()
		};

		OtherScanError(msg, self.consumed())
	}

	fn expected_min_repeats(&self, min: uint, got: uint) -> ScanError {
		OtherScanError(format!("expected at least {} repeats, got {}", min, got), self.consumed())
	}
}

#[deriving(Clone, Eq, PartialEq)]
pub struct Cursor<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> {
	slice: &'a str,
	offset: uint,
	tc: Tok,
	sp: Sp,
	cs: Cs,
}

impl<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> Cursor<'a, Tok, Sp, Cs> {
	pub fn new<'b>(s: &'b str, tc: Tok, sp: Sp, cs: Cs) -> Cursor<'b, Tok, Sp, Cs> {
		Cursor {
			slice: s,
			offset: 0,
			tc: tc,
			sp: sp,
			cs: cs,
		}
	}
}

impl<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> Show for Cursor<'a, Tok, Sp, Cs> {
	fn fmt(&self, f: &mut Formatter) -> Result<(), FormatError> {
		try!(write!(f, "Cursor<{}, {}, {}> {{ offset: {}, .. }}", self.tc, self.sp, self.cs, self.offset));
		Ok(())
	}
}

impl<'a, Tok: Tokenizer, Sp: Whitespace, Cs: CompareStrs> ScanCursor<'a> for Cursor<'a, Tok, Sp, Cs> {
	fn expect_tok(&self, s: &str) -> Result<Cursor<'a, Tok, Sp, Cs>, ScanError> {
		debug!("{}.expect_tok({})", self, s);
		match self.pop_token() {
			Some((tok, ref cur)) if self.compare_strs(s, tok) => Ok(cur.clone()),
			_ => Err(self.expected_tok(s))
		}
	}

	fn consumed(&self) -> uint {
		self.offset
	}

	fn pop_token(&self) -> Option<(&'a str, Cursor<'a, Tok, Sp, Cs>)> {
		debug!("{}.pop_token()", self);
		// First, strip out leading whitespace.  It's up to the whitespace policy to *not* strip characters it wants to turn into a token.
		let cur = self.pop_ws();

		// Next, check to see if there is a whitespace token.  This allows the space policy to do things like ignore most whitespace, but turn line breaks into explicit tokens.  Note that unlike the regular Tokenizer, the Whitespace policy is responsible for returning the str slice itself.  This is used to do things like map all whitespace to a single `" "` token.
		match self.sp.token_len(cur.tail_str()) {
			Some(end) => return Some((cur.str_slice_to(end), cur.slice_from(end))),
			None => ()
		}

		// Do not assume that empty input means we can't match a token; the token class might, for example, turn end-of-input into an explicit token.
		let tail_str = cur.tail_str();
		match self.tc.token_len(tail_str) {
			Some(end) => Some((cur.str_slice_to(end), cur.slice_from(end))),
			None => {
				// One of two things: either we have some input left and will thus return a single-character token, or there is nothing left whereby we return None.
				if cur.is_empty() {
					return None;
				} else {
					let CharRange { ch: _, next } = tail_str.char_range_at(0);
					Some((cur.str_slice_to(next), cur.slice_from(next)))
				}
			},
		}
	}

	fn pop_ws(&self) -> Cursor<'a, Tok, Sp, Cs> {
		debug!("{}.pop_ws()", self);

		self.slice_from(self.sp.strip_len(self.tail_str()))
	}

	fn slice_from(&self, from: uint) -> Cursor<'a, Tok, Sp, Cs> {
		Cursor {
			offset: ::std::cmp::min(self.slice.len(), self.offset + from),
			..self.clone()
		}
	}

	fn str_slice_to(&self, to: uint) -> &'a str {
		self.tail_str().slice_to(to)
	}

	fn str_slice_to_cur(&self, to: &Cursor<'a, Tok, Sp, Cs>) -> &'a str {
		self.slice.slice(self.offset, to.offset)
	}

	fn tail_str(&self) -> &'a str {
		self.slice.slice_from(self.offset)
	}

	fn is_empty(&self) -> bool {
		self.offset == self.slice.len()
	}

	fn compare_strs(&self, a: &str, b: &str) -> bool {
		self.cs.compare_strs(a, b)
	}
}
<|endoftext|>


<commit_before>mod host;
mod transport;

pub use host::Host;
pub use transport::Transport;

use lazy_static::lazy_static;
use regex::Regex;
use std::fmt;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::str::FromStr;

use crate::error::EndpointError;

pub type Port = u16;

/// Represents a ZMQ Endpoint.
///
/// # Examples
/// ```
/// # use zeromq::{Endpoint, Host};
/// # fn main() -> Result<(), zeromq::ZmqError> {
/// assert_eq!(
///     "tcp://example.com:4567".parse::<Endpoint>()?,
///     Endpoint::Tcp(Host::Domain("example.com".to_string()), 4567)
/// );
/// # Ok(()) }
/// ```
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Hash, Eq)]
pub enum Endpoint {
    // TODO: Add endpoints for the other transport variants
    Tcp(Host, Port),
    Ipc(PathBuf),
}

impl Endpoint {
    pub fn transport(&self) -> Transport {
        match self {
            Self::Tcp(_, _) => Transport::Tcp,
            Self::Ipc(_) => Transport::Ipc,
        }
    }

    /// Creates an `Endpoint::Tcp` from a [`SocketAddr`]
    pub fn from_tcp_sock_addr(addr: SocketAddr) -> Self {
        Endpoint::Tcp(addr.ip().into(), addr.port())
    }
}

impl FromStr for Endpoint {
    type Err = EndpointError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        lazy_static! {
            static ref TRANSPORT_REGEX: Regex = Regex::new(r"^([[:lower:]]+)://(.+)$").unwrap();
            static ref HOST_PORT_REGEX: Regex = Regex::new(r"^(.+):(\d+)$").unwrap();
        }

        let caps = TRANSPORT_REGEX
            .captures(s)
            .ok_or(EndpointError::Syntax("Could not parse transport"))?;
        let transport: &str = caps.get(1).unwrap().into();
        let transport: Transport = transport.parse()?;
        let address = caps.get(2).unwrap().as_str();

        fn extract_host_port(address: &str) -> Result<(Host, Port), EndpointError> {
            let caps = HOST_PORT_REGEX
                .captures(address)
                .ok_or(EndpointError::Syntax("Could not parse host and port"))?;
            let host = caps.get(1).unwrap().as_str();
            let port = caps.get(2).unwrap().as_str();
            let port: Port = port
                .parse()
                .map_err(|_| EndpointError::Syntax("Port must be a u16 but was out of range"))?;

            let host: Host = host.parse()?;
            Ok((host, port))
        }

        let endpoint = match transport {
            Transport::Tcp => {
                let (host, port) = extract_host_port(address)?;
                Endpoint::Tcp(host, port)
            }
            Transport::Ipc => {
                let path: PathBuf = address.to_string().into();
                Endpoint::Ipc(path)
            }
        };

        Ok(endpoint)
    }
}

impl fmt::Display for Endpoint {
    fn fmt(&self, f: &mut fmt::Formatter) -> std::result::Result<(), std::fmt::Error> {
        match self {
            Endpoint::Tcp(host, port) => {
                if let Host::Ipv6(_) = host {
                    write!(f, "tcp://[{}]:{}", host, port)
                } else {
                    write!(f, "tcp://{}:{}", host, port)
                }
            }
            Endpoint::Ipc(path) => write!(f, "ipc://{}", path.display()),
        }
    }
}

/// Represents a type that can be converted into an [`Endpoint`].
///
/// This trait is intentionally sealed to prevent implementation on third-party
/// types.
// TODO: Is sealing this trait actually necessary?
pub trait TryIntoEndpoint: Send + private::Sealed {
    /// Convert into an `Endpoint` via an owned `Self`.
    ///
    /// Enables efficient `Endpoint` -> `Endpoint` conversion, while permitting
    /// the creation of a new `Endpoint` when given types like `&str`.
    fn try_into(self) -> Result<Endpoint, EndpointError>;
}
impl TryIntoEndpoint for &str {
    fn try_into(self) -> Result<Endpoint, EndpointError> {
        self.parse()
    }
}
impl TryIntoEndpoint for Endpoint {
    fn try_into(self) -> Result<Endpoint, EndpointError> {
        Ok(self)
    }
}

impl private::Sealed for str {}
impl private::Sealed for &str {}
impl private::Sealed for Endpoint {}

mod private {
    pub trait Sealed {}
}

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

    lazy_static! {
        static ref PAIRS: Vec<(Endpoint, &'static str)> = vec![
            (
                Endpoint::Tcp(Host::Domain("www.example.com".to_string()), 1234),
                "tcp://www.example.com:1234",
            ),
            (
                Endpoint::Tcp(Host::Domain("*".to_string()), 2345),
                "tcp://*:2345",
            ),
            (
                Endpoint::Tcp(Host::Ipv4("127.0.0.1".parse().unwrap()), 8080),
                "tcp://127.0.0.1:8080",
            ),
            (
                Endpoint::Tcp(Host::Ipv6("::1".parse().unwrap()), 34567),
                "tcp://[::1]:34567",
            ),
            (
                Endpoint::Tcp(Host::Domain("i❤.ws".to_string()), 80),
                "tcp://i❤.ws:80"
            ),
            (
                Endpoint::Tcp(Host::Domain("xn--i-7iq.ws".to_string()), 80),
                "tcp://xn--i-7iq.ws:80"
            ),
            (
                Endpoint::Tcp(Host::Ipv4("127.0.0.1".parse().unwrap()), 65535),
                "tcp://127.0.0.1:65535"
            ),
            (
                Endpoint::Tcp(Host::Ipv4("127.0.0.1".parse().unwrap()), 0),
                "tcp://127.0.0.1:0"
            )
        ];
    }

    #[test]
    fn test_endpoint_display() {
        for (e, s) in PAIRS.iter() {
            assert_eq!(&format!("{}", e), s);
        }
    }

    #[test]
    fn test_endpoint_parse() {
        use std::mem::discriminant as disc;

        // Test example 1:1 pairs
        for (e, s) in PAIRS.iter() {
            assert_eq!(&s.parse::<Endpoint>().unwrap(), e);
        }

        let exact_counter_examples = vec![(
            "abc://127.0.0.1:1234",
            EndpointError::UnknownTransport("abc".to_string()),
        )];

        for (s, target_err) in exact_counter_examples {
            let parse_err = s.parse::<Endpoint>().unwrap_err();
            assert_eq!(parse_err.to_string(), target_err.to_string());
            assert_eq!(disc(&parse_err), disc(&target_err));
        }

        // Examples where we only care about matching variants, rather than contents of
        // those variants
        let inexact_counter_examples = vec![
            ("://127.0.0.1:1234", EndpointError::Syntax("")),
            ("tcp://127.0.0.1:", EndpointError::Syntax("")),
            ("tcp://:1234", EndpointError::Syntax("")),
            ("tcp://127.0.0.1", EndpointError::Syntax("")),
            ("tcp://127.0.0.1:65536", EndpointError::Syntax("")),
            ("TCP://127.0.0.1:1234", EndpointError::Syntax("")),
        ];

        for (s, target_variant) in inexact_counter_examples {
            let parse_err = s.parse::<Endpoint>().unwrap_err();
            assert_eq!(disc(&parse_err), disc(&target_variant));
        }
    }
}
<commit_msg>Added tests for ipc endpoint display<commit_after>mod host;
mod transport;

pub use host::Host;
pub use transport::Transport;

use lazy_static::lazy_static;
use regex::Regex;
use std::fmt;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::str::FromStr;

use crate::error::EndpointError;

pub type Port = u16;

/// Represents a ZMQ Endpoint.
///
/// # Examples
/// ```
/// # use zeromq::{Endpoint, Host};
/// # fn main() -> Result<(), zeromq::ZmqError> {
/// assert_eq!(
///     "tcp://example.com:4567".parse::<Endpoint>()?,
///     Endpoint::Tcp(Host::Domain("example.com".to_string()), 4567)
/// );
/// # Ok(()) }
/// ```
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Hash, Eq)]
pub enum Endpoint {
    // TODO: Add endpoints for the other transport variants
    Tcp(Host, Port),
    Ipc(PathBuf),
}

impl Endpoint {
    pub fn transport(&self) -> Transport {
        match self {
            Self::Tcp(_, _) => Transport::Tcp,
            Self::Ipc(_) => Transport::Ipc,
        }
    }

    /// Creates an `Endpoint::Tcp` from a [`SocketAddr`]
    pub fn from_tcp_sock_addr(addr: SocketAddr) -> Self {
        Endpoint::Tcp(addr.ip().into(), addr.port())
    }
}

impl FromStr for Endpoint {
    type Err = EndpointError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        lazy_static! {
            static ref TRANSPORT_REGEX: Regex = Regex::new(r"^([[:lower:]]+)://(.+)$").unwrap();
            static ref HOST_PORT_REGEX: Regex = Regex::new(r"^(.+):(\d+)$").unwrap();
        }

        let caps = TRANSPORT_REGEX
            .captures(s)
            .ok_or(EndpointError::Syntax("Could not parse transport"))?;
        let transport: &str = caps.get(1).unwrap().into();
        let transport: Transport = transport.parse()?;
        let address = caps.get(2).unwrap().as_str();

        fn extract_host_port(address: &str) -> Result<(Host, Port), EndpointError> {
            let caps = HOST_PORT_REGEX
                .captures(address)
                .ok_or(EndpointError::Syntax("Could not parse host and port"))?;
            let host = caps.get(1).unwrap().as_str();
            let port = caps.get(2).unwrap().as_str();
            let port: Port = port
                .parse()
                .map_err(|_| EndpointError::Syntax("Port must be a u16 but was out of range"))?;

            let host: Host = host.parse()?;
            Ok((host, port))
        }

        let endpoint = match transport {
            Transport::Tcp => {
                let (host, port) = extract_host_port(address)?;
                Endpoint::Tcp(host, port)
            }
            Transport::Ipc => {
                let path: PathBuf = address.to_string().into();
                Endpoint::Ipc(path)
            }
        };

        Ok(endpoint)
    }
}

impl fmt::Display for Endpoint {
    fn fmt(&self, f: &mut fmt::Formatter) -> std::result::Result<(), std::fmt::Error> {
        match self {
            Endpoint::Tcp(host, port) => {
                if let Host::Ipv6(_) = host {
                    write!(f, "tcp://[{}]:{}", host, port)
                } else {
                    write!(f, "tcp://{}:{}", host, port)
                }
            }
            Endpoint::Ipc(path) => write!(f, "ipc://{}", path.display()),
        }
    }
}

/// Represents a type that can be converted into an [`Endpoint`].
///
/// This trait is intentionally sealed to prevent implementation on third-party
/// types.
// TODO: Is sealing this trait actually necessary?
pub trait TryIntoEndpoint: Send + private::Sealed {
    /// Convert into an `Endpoint` via an owned `Self`.
    ///
    /// Enables efficient `Endpoint` -> `Endpoint` conversion, while permitting
    /// the creation of a new `Endpoint` when given types like `&str`.
    fn try_into(self) -> Result<Endpoint, EndpointError>;
}
impl TryIntoEndpoint for &str {
    fn try_into(self) -> Result<Endpoint, EndpointError> {
        self.parse()
    }
}
impl TryIntoEndpoint for Endpoint {
    fn try_into(self) -> Result<Endpoint, EndpointError> {
        Ok(self)
    }
}

impl private::Sealed for str {}
impl private::Sealed for &str {}
impl private::Sealed for Endpoint {}

mod private {
    pub trait Sealed {}
}

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

    lazy_static! {
        static ref PAIRS: Vec<(Endpoint, &'static str)> = vec![
            (Endpoint::Ipc(PathBuf::from("/tmp/asdf")), "ipc:///tmp/asdf"),
            (
                Endpoint::Ipc(PathBuf::from("my/dir_1/dir-2")),
                "ipc://my/dir_1/dir-2"
            ),
            (
                Endpoint::Ipc(PathBuf::from("@abstract/namespace")),
                "ipc://@abstract/namespace"
            ),
            (
                Endpoint::Tcp(Host::Domain("www.example.com".to_string()), 1234),
                "tcp://www.example.com:1234",
            ),
            (
                Endpoint::Tcp(Host::Domain("*".to_string()), 2345),
                "tcp://*:2345",
            ),
            (
                Endpoint::Tcp(Host::Ipv4("127.0.0.1".parse().unwrap()), 8080),
                "tcp://127.0.0.1:8080",
            ),
            (
                Endpoint::Tcp(Host::Ipv6("::1".parse().unwrap()), 34567),
                "tcp://[::1]:34567",
            ),
            (
                Endpoint::Tcp(Host::Domain("i❤.ws".to_string()), 80),
                "tcp://i❤.ws:80"
            ),
            (
                Endpoint::Tcp(Host::Domain("xn--i-7iq.ws".to_string()), 80),
                "tcp://xn--i-7iq.ws:80"
            ),
            (
                Endpoint::Tcp(Host::Ipv4("127.0.0.1".parse().unwrap()), 65535),
                "tcp://127.0.0.1:65535"
            ),
            (
                Endpoint::Tcp(Host::Ipv4("127.0.0.1".parse().unwrap()), 0),
                "tcp://127.0.0.1:0"
            )
        ];
    }

    #[test]
    fn test_endpoint_display() {
        for (e, s) in PAIRS.iter() {
            assert_eq!(&format!("{}", e), s);
        }
    }

    #[test]
    fn test_endpoint_parse() {
        use std::mem::discriminant as disc;

        // Test example 1:1 pairs
        for (e, s) in PAIRS.iter() {
            assert_eq!(&s.parse::<Endpoint>().unwrap(), e);
        }

        let exact_counter_examples = vec![(
            "abc://127.0.0.1:1234",
            EndpointError::UnknownTransport("abc".to_string()),
        )];

        for (s, target_err) in exact_counter_examples {
            let parse_err = s.parse::<Endpoint>().unwrap_err();
            assert_eq!(parse_err.to_string(), target_err.to_string());
            assert_eq!(disc(&parse_err), disc(&target_err));
        }

        // Examples where we only care about matching variants, rather than contents of
        // those variants
        let inexact_counter_examples = vec![
            ("://127.0.0.1:1234", EndpointError::Syntax("")),
            ("tcp://127.0.0.1:", EndpointError::Syntax("")),
            ("tcp://:1234", EndpointError::Syntax("")),
            ("tcp://127.0.0.1", EndpointError::Syntax("")),
            ("tcp://127.0.0.1:65536", EndpointError::Syntax("")),
            ("TCP://127.0.0.1:1234", EndpointError::Syntax("")),
        ];

        for (s, target_variant) in inexact_counter_examples {
            let parse_err = s.parse::<Endpoint>().unwrap_err();
            assert_eq!(disc(&parse_err), disc(&target_variant));
        }
    }
}
<|endoftext|>


<commit_before>use crate::value::AMQPValue;

use std::{
    collections::{BTreeMap, btree_map},
    borrow, fmt, str,
};

use serde::{Deserialize, Serialize};

/// Enumeration referencing all the available AMQP types
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
pub enum AMQPType {
    /// A bool
    Boolean,
    /// An i8
    ShortShortInt,
    /// A u8
    ShortShortUInt,
    /// An i16
    ShortInt,
    /// A u16
    ShortUInt,
    /// An i32
    LongInt,
    /// A u32
    LongUInt,
    /// An i64
    LongLongInt,
    /// A u64
    LongLongUInt,
    /// An f32
    Float,
    /// An f64
    Double,
    /// A decimal value represented by a scale and a value
    DecimalValue,
    /// Deprecated, a String
    ShortString,
    /// A String
    LongString,
    /// An array of AMQPValue
    FieldArray,
    /// A timestamp (u32)
    Timestamp,
    /// A Map<String, AMQPValue>
    FieldTable,
    /// An array of bytes, RabbitMQ specific
    ByteArray, /* ByteArray is specific to RabbitMQ */
    /// No value
    Void,
}

impl AMQPType {
    /// Get the AMQPType corresponding to the given id.
    /// We don't strictly follow the spec here but rather the RabbitMQ implementation
    /// 's' means ShortInt (like 'U') instead of ShortString
    /// 'l' and 'L' both mean LongLongInt (no LongLongUInt)
    pub fn from_id(id: char) -> Option<AMQPType> {
        match id {
            't' => Some(AMQPType::Boolean),
            'b' => Some(AMQPType::ShortShortInt),
            'B' => Some(AMQPType::ShortShortUInt),
            /* Specs says 'U', RabbitMQ says 's' (which means ShortString in specs) */
            's' |
            'U' => Some(AMQPType::ShortInt),
            'u' => Some(AMQPType::ShortUInt),
            'I' => Some(AMQPType::LongInt),
            'i' => Some(AMQPType::LongUInt),
            /* RabbitMQ treats both 'l' and 'L' as LongLongInt and ignores LongLongUInt */
            'L' |
            'l' => Some(AMQPType::LongLongInt),
            'f' => Some(AMQPType::Float),
            'd' => Some(AMQPType::Double),
            'D' => Some(AMQPType::DecimalValue),
            'S' => Some(AMQPType::LongString),
            'A' => Some(AMQPType::FieldArray),
            'T' => Some(AMQPType::Timestamp),
            'F' => Some(AMQPType::FieldTable),
            'x' => Some(AMQPType::ByteArray),
            'V' => Some(AMQPType::Void),
            _   => None,
        }
    }

    /// Get the id from an AMQPType
    /// We don't strictly follow the spec here but rather the RabbitMQ implementation
    /// ShortString doesn't have an id, we return '_' instead
    /// ShortInt is supposed to be 'U' but we use 's'
    /// LongLongUInt is supposed to be 'L' but we return 'l' as LongLongInt
    pub fn get_id(self) -> char {
        match self {
            AMQPType::Boolean        => 't',
            AMQPType::ShortShortInt  => 'b',
            AMQPType::ShortShortUInt => 'B',
            /* Specs says 'U', RabbitMQ says 's' (which means ShortString in specs) */
            AMQPType::ShortInt       => 's',
            AMQPType::ShortUInt      => 'u',
            AMQPType::LongInt        => 'I',
            AMQPType::LongUInt       => 'i',
            /* RabbitMQ treats both 'l' and 'L' as LongLongInt and ignores LongLongUInt */
            AMQPType::LongLongInt    |
            AMQPType::LongLongUInt   => 'l',
            AMQPType::Float          => 'f',
            AMQPType::Double         => 'd',
            AMQPType::DecimalValue   => 'D',
            /* ShortString only exists for internal usage, we shouldn't ever have to use this */
            AMQPType::ShortString    => '_',
            AMQPType::LongString     => 'S',
            AMQPType::FieldArray     => 'A',
            AMQPType::Timestamp      => 'T',
            AMQPType::FieldTable     => 'F',
            AMQPType::ByteArray      => 'x',
            AMQPType::Void           => 'V',
        }
    }
}

impl fmt::Display for AMQPType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

/// A bool
pub type Boolean        = bool;
/// An i8
pub type ShortShortInt  = i8;
/// A u8
pub type ShortShortUInt = u8;
/// An i16
pub type ShortInt       = i16;
/// A u16
pub type ShortUInt      = u16;
/// An i32
pub type LongInt        = i32;
/// A u32
pub type LongUInt       = u32;
/// An i64
pub type LongLongInt    = i64;
/// A u64
pub type LongLongUInt   = u64;
/// A f32
pub type Float          = f32;
/// A f64
pub type Double         = f64;
/// A timestamp (u32)
pub type Timestamp      = LongLongUInt;
/// No value
pub type Void           = ();

/// A String (deprecated)
#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
pub struct ShortString(String);
/// A String
#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
pub struct LongString(String);
/// An array of AMQPValue
#[derive(Clone, Debug, Default, PartialEq, Deserialize, Serialize)]
pub struct FieldArray(Vec<AMQPValue>);
/// A Map<String, AMQPValue>
#[derive(Clone, Debug, Default, PartialEq, Deserialize, Serialize)]
pub struct FieldTable(BTreeMap<ShortString, AMQPValue>);
/// An array of bytes (RabbitMQ specific)
#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Deserialize, Serialize)]
pub struct ByteArray(Vec<u8>);

/// A Decimal value composed of a scale and a value
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Deserialize, Serialize)]
pub struct DecimalValue {
    /// The scale of the value
    pub scale: ShortShortUInt,
    /// The actual value
    pub value: LongUInt,
}

/// A Reference pointing to a ShortString
pub struct ShortStringRef<'a>(pub &'a str);
/// A Reference pointing to a LongString
pub struct LongStringRef<'a>(pub &'a str);

impl<'a> ShortString {
    /// Get a reference to a ShortString
    pub fn as_ref(&'a self) -> ShortStringRef<'a> {
        ShortStringRef(&self.0)
    }

    /// Get a reference to a LongString as &str
    pub fn as_str(&'a self) -> &'a str {
        self.0.as_str()
    }

    /// Splits a string slice by whitespace.
    pub fn split_whitespace(&'a self) -> str::SplitWhitespace<'a> {
        self.0.split_whitespace()
    }
}

impl From<String> for ShortString {
    fn from(s: String) -> Self {
        Self(s)
    }
}

impl From<&str> for ShortString {
    fn from(s: &str) -> Self {
        s.to_owned().into()
    }
}

impl borrow::Borrow<str> for ShortString {
    fn borrow(&self) -> &str {
        self.0.borrow()
    }
}

impl fmt::Display for ShortString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<'a> ShortStringRef<'a> {
    pub(crate) fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

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

impl Default for ShortStringRef<'static> {
    fn default() -> Self {
        Self("")
    }
}

impl<'a> From<&'a str> for ShortStringRef<'a> {
    fn from(s: &'a str) -> Self {
        Self(s)
    }
}

impl<'a> LongString {
    /// Get a reference to a LongString
    pub fn as_ref(&'a self) -> LongStringRef<'a> {
        LongStringRef(&self.0)
    }

    /// Get a reference to a LongString as &str
    pub fn as_str(&'a self) -> &'a str {
        self.0.as_str()
    }

    /// Splits a string slice by whitespace.
    pub fn split_whitespace(&'a self) -> str::SplitWhitespace<'a> {
        self.0.split_whitespace()
    }
}

impl From<String> for LongString {
    fn from(s: String) -> Self {
        Self(s)
    }
}

impl From<&str> for LongString {
    fn from(s: &str) -> Self {
        s.to_owned().into()
    }
}

impl borrow::Borrow<str> for LongString {
    fn borrow(&self) -> &str {
        self.0.borrow()
    }
}

impl fmt::Display for LongString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<'a> LongStringRef<'a> {
    pub(crate) fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

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

impl Default for LongStringRef<'static> {
    fn default() -> Self {
        Self("")
    }
}

impl<'a> From<&'a str> for LongStringRef<'a> {
    fn from(s: &'a str) -> Self {
        Self(s)
    }
}

impl FieldArray {
    pub(crate) fn as_slice(&self) -> &[AMQPValue] {
        self.0.as_slice()
    }

    pub(crate) fn push(&mut self, v: AMQPValue) {
        self.0.push(v);
    }
}

impl From<Vec<AMQPValue>> for FieldArray {
    fn from(v: Vec<AMQPValue>) -> Self {
        Self(v)
    }
}

impl FieldTable {
    /// Insert a new entry in the table
    pub fn insert(&mut self, k: ShortString, v: AMQPValue) {
        self.0.insert(k, v);
    }

    /// Check whether the table contains the given key
    pub fn contains_key(&self, k: &str) -> bool {
        self.0.contains_key(k)
    }

    /// Access the inner BTreeMap to perform lookups
    pub fn inner(&self) -> &BTreeMap<ShortString, AMQPValue> {
        &self.0
    }
}

impl<'a> IntoIterator for &'a FieldTable {
    type Item = (&'a ShortString, &'a AMQPValue);
    type IntoIter = btree_map::Iter<'a, ShortString, AMQPValue>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl From<BTreeMap<ShortString, AMQPValue>> for FieldTable {
    fn from(m: BTreeMap<ShortString, AMQPValue>) -> Self {
        Self(m)
    }
}

impl ByteArray {
    pub(crate) fn as_slice(&self) -> &[u8] {
        self.0.as_slice()
    }

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

impl From<Vec<u8>> for ByteArray {
    fn from(v: Vec<u8>) -> Self {
        Self(v)
    }
}

impl From<&[u8]> for ByteArray {
    fn from(v: &[u8]) -> Self {
        Self(v.to_vec())
    }
}

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

    #[test]
    fn test_type_from_id() {
        assert_eq!(AMQPType::from_id('T'), Some(AMQPType::Timestamp));
        assert_eq!(AMQPType::from_id('S'), Some(AMQPType::LongString));
        assert_eq!(AMQPType::from_id('s'), Some(AMQPType::ShortInt));
        assert_eq!(AMQPType::from_id('U'), Some(AMQPType::ShortInt));
        assert_eq!(AMQPType::from_id('l'), Some(AMQPType::LongLongInt));
        assert_eq!(AMQPType::from_id('z'), None);
    }

    #[test]
    fn test_type_get_id() {
        assert_eq!(AMQPType::LongLongInt.get_id(),  'l');
        assert_eq!(AMQPType::LongLongUInt.get_id(), 'l');
        assert_eq!(AMQPType::ShortString.get_id(),  '_');
    }

    #[test]
    fn test_type_to_string() {
        assert_eq!(AMQPType::Boolean.to_string(), "Boolean");
        assert_eq!(AMQPType::Void.to_string(),    "Void");
    }
}
<commit_msg>types: make some methods public<commit_after>use crate::value::AMQPValue;

use std::{
    collections::{BTreeMap, btree_map},
    borrow, fmt, str,
};

use serde::{Deserialize, Serialize};

/// Enumeration referencing all the available AMQP types
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
pub enum AMQPType {
    /// A bool
    Boolean,
    /// An i8
    ShortShortInt,
    /// A u8
    ShortShortUInt,
    /// An i16
    ShortInt,
    /// A u16
    ShortUInt,
    /// An i32
    LongInt,
    /// A u32
    LongUInt,
    /// An i64
    LongLongInt,
    /// A u64
    LongLongUInt,
    /// An f32
    Float,
    /// An f64
    Double,
    /// A decimal value represented by a scale and a value
    DecimalValue,
    /// Deprecated, a String
    ShortString,
    /// A String
    LongString,
    /// An array of AMQPValue
    FieldArray,
    /// A timestamp (u32)
    Timestamp,
    /// A Map<String, AMQPValue>
    FieldTable,
    /// An array of bytes, RabbitMQ specific
    ByteArray, /* ByteArray is specific to RabbitMQ */
    /// No value
    Void,
}

impl AMQPType {
    /// Get the AMQPType corresponding to the given id.
    /// We don't strictly follow the spec here but rather the RabbitMQ implementation
    /// 's' means ShortInt (like 'U') instead of ShortString
    /// 'l' and 'L' both mean LongLongInt (no LongLongUInt)
    pub fn from_id(id: char) -> Option<AMQPType> {
        match id {
            't' => Some(AMQPType::Boolean),
            'b' => Some(AMQPType::ShortShortInt),
            'B' => Some(AMQPType::ShortShortUInt),
            /* Specs says 'U', RabbitMQ says 's' (which means ShortString in specs) */
            's' |
            'U' => Some(AMQPType::ShortInt),
            'u' => Some(AMQPType::ShortUInt),
            'I' => Some(AMQPType::LongInt),
            'i' => Some(AMQPType::LongUInt),
            /* RabbitMQ treats both 'l' and 'L' as LongLongInt and ignores LongLongUInt */
            'L' |
            'l' => Some(AMQPType::LongLongInt),
            'f' => Some(AMQPType::Float),
            'd' => Some(AMQPType::Double),
            'D' => Some(AMQPType::DecimalValue),
            'S' => Some(AMQPType::LongString),
            'A' => Some(AMQPType::FieldArray),
            'T' => Some(AMQPType::Timestamp),
            'F' => Some(AMQPType::FieldTable),
            'x' => Some(AMQPType::ByteArray),
            'V' => Some(AMQPType::Void),
            _   => None,
        }
    }

    /// Get the id from an AMQPType
    /// We don't strictly follow the spec here but rather the RabbitMQ implementation
    /// ShortString doesn't have an id, we return '_' instead
    /// ShortInt is supposed to be 'U' but we use 's'
    /// LongLongUInt is supposed to be 'L' but we return 'l' as LongLongInt
    pub fn get_id(self) -> char {
        match self {
            AMQPType::Boolean        => 't',
            AMQPType::ShortShortInt  => 'b',
            AMQPType::ShortShortUInt => 'B',
            /* Specs says 'U', RabbitMQ says 's' (which means ShortString in specs) */
            AMQPType::ShortInt       => 's',
            AMQPType::ShortUInt      => 'u',
            AMQPType::LongInt        => 'I',
            AMQPType::LongUInt       => 'i',
            /* RabbitMQ treats both 'l' and 'L' as LongLongInt and ignores LongLongUInt */
            AMQPType::LongLongInt    |
            AMQPType::LongLongUInt   => 'l',
            AMQPType::Float          => 'f',
            AMQPType::Double         => 'd',
            AMQPType::DecimalValue   => 'D',
            /* ShortString only exists for internal usage, we shouldn't ever have to use this */
            AMQPType::ShortString    => '_',
            AMQPType::LongString     => 'S',
            AMQPType::FieldArray     => 'A',
            AMQPType::Timestamp      => 'T',
            AMQPType::FieldTable     => 'F',
            AMQPType::ByteArray      => 'x',
            AMQPType::Void           => 'V',
        }
    }
}

impl fmt::Display for AMQPType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

/// A bool
pub type Boolean        = bool;
/// An i8
pub type ShortShortInt  = i8;
/// A u8
pub type ShortShortUInt = u8;
/// An i16
pub type ShortInt       = i16;
/// A u16
pub type ShortUInt      = u16;
/// An i32
pub type LongInt        = i32;
/// A u32
pub type LongUInt       = u32;
/// An i64
pub type LongLongInt    = i64;
/// A u64
pub type LongLongUInt   = u64;
/// A f32
pub type Float          = f32;
/// A f64
pub type Double         = f64;
/// A timestamp (u32)
pub type Timestamp      = LongLongUInt;
/// No value
pub type Void           = ();

/// A String (deprecated)
#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
pub struct ShortString(String);
/// A String
#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
pub struct LongString(String);
/// An array of AMQPValue
#[derive(Clone, Debug, Default, PartialEq, Deserialize, Serialize)]
pub struct FieldArray(Vec<AMQPValue>);
/// A Map<String, AMQPValue>
#[derive(Clone, Debug, Default, PartialEq, Deserialize, Serialize)]
pub struct FieldTable(BTreeMap<ShortString, AMQPValue>);
/// An array of bytes (RabbitMQ specific)
#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Deserialize, Serialize)]
pub struct ByteArray(Vec<u8>);

/// A Decimal value composed of a scale and a value
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Deserialize, Serialize)]
pub struct DecimalValue {
    /// The scale of the value
    pub scale: ShortShortUInt,
    /// The actual value
    pub value: LongUInt,
}

/// A Reference pointing to a ShortString
pub struct ShortStringRef<'a>(pub &'a str);
/// A Reference pointing to a LongString
pub struct LongStringRef<'a>(pub &'a str);

impl<'a> ShortString {
    /// Get a reference to a ShortString
    pub fn as_ref(&'a self) -> ShortStringRef<'a> {
        ShortStringRef(&self.0)
    }

    /// Get a reference to a LongString as &str
    pub fn as_str(&'a self) -> &'a str {
        self.0.as_str()
    }

    /// Splits a string slice by whitespace.
    pub fn split_whitespace(&'a self) -> str::SplitWhitespace<'a> {
        self.0.split_whitespace()
    }
}

impl From<String> for ShortString {
    fn from(s: String) -> Self {
        Self(s)
    }
}

impl From<&str> for ShortString {
    fn from(s: &str) -> Self {
        s.to_owned().into()
    }
}

impl borrow::Borrow<str> for ShortString {
    fn borrow(&self) -> &str {
        self.0.borrow()
    }
}

impl fmt::Display for ShortString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<'a> ShortStringRef<'a> {
    pub(crate) fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

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

impl Default for ShortStringRef<'static> {
    fn default() -> Self {
        Self("")
    }
}

impl<'a> From<&'a str> for ShortStringRef<'a> {
    fn from(s: &'a str) -> Self {
        Self(s)
    }
}

impl<'a> LongString {
    /// Get a reference to a LongString
    pub fn as_ref(&'a self) -> LongStringRef<'a> {
        LongStringRef(&self.0)
    }

    /// Get a reference to a LongString as &str
    pub fn as_str(&'a self) -> &'a str {
        self.0.as_str()
    }

    /// Splits a string slice by whitespace.
    pub fn split_whitespace(&'a self) -> str::SplitWhitespace<'a> {
        self.0.split_whitespace()
    }
}

impl From<String> for LongString {
    fn from(s: String) -> Self {
        Self(s)
    }
}

impl From<&str> for LongString {
    fn from(s: &str) -> Self {
        s.to_owned().into()
    }
}

impl borrow::Borrow<str> for LongString {
    fn borrow(&self) -> &str {
        self.0.borrow()
    }
}

impl fmt::Display for LongString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<'a> LongStringRef<'a> {
    pub(crate) fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

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

impl Default for LongStringRef<'static> {
    fn default() -> Self {
        Self("")
    }
}

impl<'a> From<&'a str> for LongStringRef<'a> {
    fn from(s: &'a str) -> Self {
        Self(s)
    }
}

impl FieldArray {
    /// Get the inner values as a slice
    pub fn as_slice(&self) -> &[AMQPValue] {
        self.0.as_slice()
    }

    /// Add an item to the array
    pub fn push(&mut self, v: AMQPValue) {
        self.0.push(v);
    }
}

impl From<Vec<AMQPValue>> for FieldArray {
    fn from(v: Vec<AMQPValue>) -> Self {
        Self(v)
    }
}

impl FieldTable {
    /// Insert a new entry in the table
    pub fn insert(&mut self, k: ShortString, v: AMQPValue) {
        self.0.insert(k, v);
    }

    /// Check whether the table contains the given key
    pub fn contains_key(&self, k: &str) -> bool {
        self.0.contains_key(k)
    }

    /// Access the inner BTreeMap to perform lookups
    pub fn inner(&self) -> &BTreeMap<ShortString, AMQPValue> {
        &self.0
    }
}

impl<'a> IntoIterator for &'a FieldTable {
    type Item = (&'a ShortString, &'a AMQPValue);
    type IntoIter = btree_map::Iter<'a, ShortString, AMQPValue>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl From<BTreeMap<ShortString, AMQPValue>> for FieldTable {
    fn from(m: BTreeMap<ShortString, AMQPValue>) -> Self {
        Self(m)
    }
}

impl ByteArray {
    /// Get the inner bytes array as slice
    pub fn as_slice(&self) -> &[u8] {
        self.0.as_slice()
    }

    /// Get the length of the inner bytes array
    pub fn len(&self) -> usize {
        self.0.len()
    }
}

impl From<Vec<u8>> for ByteArray {
    fn from(v: Vec<u8>) -> Self {
        Self(v)
    }
}

impl From<&[u8]> for ByteArray {
    fn from(v: &[u8]) -> Self {
        Self(v.to_vec())
    }
}

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

    #[test]
    fn test_type_from_id() {
        assert_eq!(AMQPType::from_id('T'), Some(AMQPType::Timestamp));
        assert_eq!(AMQPType::from_id('S'), Some(AMQPType::LongString));
        assert_eq!(AMQPType::from_id('s'), Some(AMQPType::ShortInt));
        assert_eq!(AMQPType::from_id('U'), Some(AMQPType::ShortInt));
        assert_eq!(AMQPType::from_id('l'), Some(AMQPType::LongLongInt));
        assert_eq!(AMQPType::from_id('z'), None);
    }

    #[test]
    fn test_type_get_id() {
        assert_eq!(AMQPType::LongLongInt.get_id(),  'l');
        assert_eq!(AMQPType::LongLongUInt.get_id(), 'l');
        assert_eq!(AMQPType::ShortString.get_id(),  '_');
    }

    #[test]
    fn test_type_to_string() {
        assert_eq!(AMQPType::Boolean.to_string(), "Boolean");
        assert_eq!(AMQPType::Void.to_string(),    "Void");
    }
}
<|endoftext|>


<commit_before>use cssparser::Parser;
use markup5ever::{expanded_name, local_name, namespace_url, ns, QualName};
use nalgebra::{DMatrix, Dynamic, VecStorage};

use crate::attributes::Attributes;
use crate::document::AcquiredNodes;
use crate::drawing_ctx::DrawingCtx;
use crate::element::{ElementResult, SetAttributes};
use crate::error::*;
use crate::node::Node;
use crate::number_list::{NumberList, NumberListLength};
use crate::parsers::{NumberOptionalNumber, Parse, ParseValue};
use crate::rect::IRect;
use crate::surface_utils::{
    iterators::{PixelRectangle, Pixels},
    shared_surface::ExclusiveImageSurface,
    EdgeMode, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;

use super::context::{FilterContext, FilterOutput, FilterResult};
use super::{FilterEffect, FilterError, PrimitiveWithInput};

/// The `feConvolveMatrix` filter primitive.
pub struct FeConvolveMatrix {
    base: PrimitiveWithInput,
    order: (u32, u32),
    kernel_matrix: Option<DMatrix<f64>>,
    divisor: Option<f64>,
    bias: f64,
    target_x: Option<u32>,
    target_y: Option<u32>,
    edge_mode: EdgeMode,
    kernel_unit_length: Option<(f64, f64)>,
    preserve_alpha: bool,
}

impl Default for FeConvolveMatrix {
    /// Constructs a new `ConvolveMatrix` with empty properties.
    #[inline]
    fn default() -> FeConvolveMatrix {
        FeConvolveMatrix {
            base: PrimitiveWithInput::new::<Self>(),
            order: (3, 3),
            kernel_matrix: None,
            divisor: None,
            bias: 0.0,
            target_x: None,
            target_y: None,
            edge_mode: EdgeMode::Duplicate,
            kernel_unit_length: None,
            preserve_alpha: false,
        }
    }
}

impl SetAttributes for FeConvolveMatrix {
    fn set_attributes(&mut self, attrs: &Attributes) -> ElementResult {
        self.base.set_attributes(attrs)?;

        for (attr, value) in attrs.iter() {
            match attr.expanded() {
                expanded_name!("", "order") => {
                    let NumberOptionalNumber(x, y) =
                        attr.parse_and_validate(value, |v: NumberOptionalNumber<i32>| {
                            if v.0 > 0 && v.1 > 0 {
                                Ok(v)
                            } else {
                                Err(ValueErrorKind::value_error("values must be greater than 0"))
                            }
                        })?;
                    self.order = (x as u32, y as u32);
                }
                expanded_name!("", "divisor") => {
                    self.divisor = Some(attr.parse_and_validate(value, |x| {
                        if x != 0.0 {
                            Ok(x)
                        } else {
                            Err(ValueErrorKind::value_error("divisor cannot be equal to 0"))
                        }
                    })?)
                }
                expanded_name!("", "bias") => self.bias = attr.parse(value)?,
                expanded_name!("", "edgeMode") => self.edge_mode = attr.parse(value)?,
                expanded_name!("", "kernelUnitLength") => {
                    let NumberOptionalNumber(x, y) =
                        attr.parse_and_validate(value, |v: NumberOptionalNumber<f64>| {
                            if v.0 > 0.0 && v.1 > 0.0 {
                                Ok(v)
                            } else {
                                Err(ValueErrorKind::value_error(
                                    "kernelUnitLength can't be less or equal to zero",
                                ))
                            }
                        })?;

                    self.kernel_unit_length = Some((x, y))
                }
                expanded_name!("", "preserveAlpha") => self.preserve_alpha = attr.parse(value)?,

                _ => (),
            }
        }

        // target_x and target_y depend on order.
        for (attr, value) in attrs.iter() {
            match attr.expanded() {
                expanded_name!("", "targetX") => {
                    self.target_x = {
                        let v = attr.parse_and_validate(value, |v: i32| {
                            if v >= 0 && v < self.order.0 as i32 {
                                Ok(v)
                            } else {
                                Err(ValueErrorKind::value_error(
                                    "targetX must be greater or equal to zero and less than orderX",
                                ))
                            }
                        })?;
                        Some(v as u32)
                    }
                }
                expanded_name!("", "targetY") => {
                    self.target_y = {
                        let v = attr.parse_and_validate(value, |v: i32| {
                            if v >= 0 && v < self.order.1 as i32 {
                                Ok(v)
                            } else {
                                Err(ValueErrorKind::value_error(
                                    "targetY must be greater or equal to zero and less than orderY",
                                ))
                            }
                        })?;
                        Some(v as u32)
                    }
                }
                _ => (),
            }
        }

        // Default values for target_x and target_y.
        if self.target_x.is_none() {
            self.target_x = Some(self.order.0 / 2);
        }
        if self.target_y.is_none() {
            self.target_y = Some(self.order.1 / 2);
        }

        // Finally, parse the kernel matrix.
        for (attr, value) in attrs
            .iter()
            .filter(|(attr, _)| attr.expanded() == expanded_name!("", "kernelMatrix"))
        {
            self.kernel_matrix = Some({
                let number_of_elements = self.order.0 as usize * self.order.1 as usize;

                // #352: Parse as an unbounded list rather than exact length to prevent aborts due
                //       to huge allocation attempts by underlying Vec::with_capacity().
                let NumberList(v) = NumberList::parse_str(value, NumberListLength::Unbounded)
                    .attribute(attr.clone())?;

                if v.len() != number_of_elements {
                    return Err(ValueErrorKind::value_error(&format!(
                        "incorrect number of elements: expected {}",
                        number_of_elements
                    )))
                    .attribute(attr);
                }

                DMatrix::from_data(VecStorage::new(
                    Dynamic::new(self.order.1 as usize),
                    Dynamic::new(self.order.0 as usize),
                    v,
                ))
            });
        }

        // kernel_matrix must have been specified.
        if self.kernel_matrix.is_none() {
            return Err(ValueErrorKind::value_error("the value must be set"))
                .attribute(QualName::new(None, ns!(svg), local_name!("kernelMatrix")));
        }

        // Default value for the divisor.
        if self.divisor.is_none() {
            self.divisor = Some(self.kernel_matrix.as_ref().unwrap().iter().sum());

            if self.divisor.unwrap() == 0.0 {
                self.divisor = Some(1.0);
            }
        }

        Ok(())
    }
}

impl FilterEffect for FeConvolveMatrix {
    fn render(
        &self,
        node: &Node,
        ctx: &FilterContext,
        acquired_nodes: &mut AcquiredNodes<'_>,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError> {
        #![allow(clippy::many_single_char_names)]

        let input = self.base.get_input(ctx, acquired_nodes, draw_ctx)?;
        let mut bounds = self
            .base
            .get_bounds(ctx, node.parent().as_ref())?
            .add_input(&input)
            .into_irect(draw_ctx);
        let original_bounds = bounds;

        let mut input_surface = if self.preserve_alpha {
            // preserve_alpha means we need to premultiply and unpremultiply the values.
            input.surface().unpremultiply(bounds)?
        } else {
            input.surface().clone()
        };

        let scale = self
            .kernel_unit_length
            .map(|(dx, dy)| ctx.paffine().transform_distance(dx, dy));

        if let Some((ox, oy)) = scale {
            // Scale the input surface to match kernel_unit_length.
            let (new_surface, new_bounds) = input_surface.scale(bounds, 1.0 / ox, 1.0 / oy)?;

            input_surface = new_surface;
            bounds = new_bounds;
        }

        let matrix = self.kernel_matrix.as_ref().unwrap();

        let mut surface = ExclusiveImageSurface::new(
            input_surface.width(),
            input_surface.height(),
            input.surface().surface_type(),
        )?;

        surface.modify(&mut |data, stride| {
            for (x, y, pixel) in Pixels::within(&input_surface, bounds) {
                // Compute the convolution rectangle bounds.
                let kernel_bounds = IRect::new(
                    x as i32 - self.target_x.unwrap() as i32,
                    y as i32 - self.target_y.unwrap() as i32,
                    x as i32 - self.target_x.unwrap() as i32 + self.order.0 as i32,
                    y as i32 - self.target_y.unwrap() as i32 + self.order.1 as i32,
                );

                // Do the convolution.
                let mut r = 0.0;
                let mut g = 0.0;
                let mut b = 0.0;
                let mut a = 0.0;

                for (x, y, pixel) in
                    PixelRectangle::within(&input_surface, bounds, kernel_bounds, self.edge_mode)
                {
                    let kernel_x = (kernel_bounds.x1 - x - 1) as usize;
                    let kernel_y = (kernel_bounds.y1 - y - 1) as usize;

                    r += f64::from(pixel.r) / 255.0 * matrix[(kernel_y, kernel_x)];
                    g += f64::from(pixel.g) / 255.0 * matrix[(kernel_y, kernel_x)];
                    b += f64::from(pixel.b) / 255.0 * matrix[(kernel_y, kernel_x)];

                    if !self.preserve_alpha {
                        a += f64::from(pixel.a) / 255.0 * matrix[(kernel_y, kernel_x)];
                    }
                }

                // If preserve_alpha is true, set a to the source alpha value.
                if self.preserve_alpha {
                    a = f64::from(pixel.a) / 255.0;
                } else {
                    a = a / self.divisor.unwrap() + self.bias;
                }

                let clamped_a = clamp(a, 0.0, 1.0);

                let compute = |x| {
                    let x = x / self.divisor.unwrap() + self.bias * a;

                    let x = if self.preserve_alpha {
                        // Premultiply the output value.
                        clamp(x, 0.0, 1.0) * clamped_a
                    } else {
                        clamp(x, 0.0, clamped_a)
                    };

                    ((x * 255.0) + 0.5) as u8
                };

                let output_pixel = Pixel {
                    r: compute(r),
                    g: compute(g),
                    b: compute(b),
                    a: ((clamped_a * 255.0) + 0.5) as u8,
                };

                data.set_pixel(stride, output_pixel, x, y);
            }
        });

        let mut surface = surface.share()?;

        if let Some((ox, oy)) = scale {
            // Scale the output surface back.
            surface = surface.scale_to(
                ctx.source_graphic().width(),
                ctx.source_graphic().height(),
                original_bounds,
                ox,
                oy,
            )?;

            bounds = original_bounds;
        }

        Ok(FilterResult {
            name: self.base.result.clone(),
            output: FilterOutput { surface, bounds },
        })
    }

    #[inline]
    fn is_affected_by_color_interpolation_filters(&self) -> bool {
        true
    }
}

impl Parse for EdgeMode {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<Self, ParseError<'i>> {
        Ok(parse_identifiers!(
            parser,
            "duplicate" => EdgeMode::Duplicate,
            "wrap" => EdgeMode::Wrap,
            "none" => EdgeMode::None,
        )?)
    }
}

// Used for the preserveAlpha attribute
impl Parse for bool {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<Self, ParseError<'i>> {
        Ok(parse_identifiers!(
            parser,
            "false" => false,
            "true" => true,
        )?)
    }
}
<commit_msg>convolve_matrix: use NonNegative and NonZero for parsing<commit_after>use cssparser::Parser;
use markup5ever::{expanded_name, local_name, namespace_url, ns, QualName};
use nalgebra::{DMatrix, Dynamic, VecStorage};

use crate::attributes::Attributes;
use crate::document::AcquiredNodes;
use crate::drawing_ctx::DrawingCtx;
use crate::element::{ElementResult, SetAttributes};
use crate::error::*;
use crate::node::Node;
use crate::number_list::{NumberList, NumberListLength};
use crate::parsers::{NonNegative, NonZero, NumberOptionalNumber, Parse, ParseValue};
use crate::rect::IRect;
use crate::surface_utils::{
    iterators::{PixelRectangle, Pixels},
    shared_surface::ExclusiveImageSurface,
    EdgeMode, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;

use super::context::{FilterContext, FilterOutput, FilterResult};
use super::{FilterEffect, FilterError, PrimitiveWithInput};

/// The `feConvolveMatrix` filter primitive.
pub struct FeConvolveMatrix {
    base: PrimitiveWithInput,
    order: (u32, u32),
    kernel_matrix: Option<DMatrix<f64>>,
    divisor: Option<f64>,
    bias: f64,
    target_x: Option<u32>,
    target_y: Option<u32>,
    edge_mode: EdgeMode,
    kernel_unit_length: Option<(f64, f64)>,
    preserve_alpha: bool,
}

impl Default for FeConvolveMatrix {
    /// Constructs a new `ConvolveMatrix` with empty properties.
    #[inline]
    fn default() -> FeConvolveMatrix {
        FeConvolveMatrix {
            base: PrimitiveWithInput::new::<Self>(),
            order: (3, 3),
            kernel_matrix: None,
            divisor: None,
            bias: 0.0,
            target_x: None,
            target_y: None,
            edge_mode: EdgeMode::Duplicate,
            kernel_unit_length: None,
            preserve_alpha: false,
        }
    }
}

impl SetAttributes for FeConvolveMatrix {
    fn set_attributes(&mut self, attrs: &Attributes) -> ElementResult {
        self.base.set_attributes(attrs)?;

        for (attr, value) in attrs.iter() {
            match attr.expanded() {
                expanded_name!("", "order") => {
                    let NumberOptionalNumber(NonNegative(x), NonNegative(y)) = attr.parse(value)?;
                    self.order = (x as u32, y as u32);
                }
                expanded_name!("", "divisor") => {
                    let NonZero(d) = attr.parse(value)?;
                    self.divisor = Some(d);
                }
                expanded_name!("", "bias") => self.bias = attr.parse(value)?,
                expanded_name!("", "edgeMode") => self.edge_mode = attr.parse(value)?,
                expanded_name!("", "kernelUnitLength") => {
                    let NumberOptionalNumber(NonNegative(x), NonNegative(y)) = attr.parse(value)?;
                    self.kernel_unit_length = Some((x, y))
                }
                expanded_name!("", "preserveAlpha") => self.preserve_alpha = attr.parse(value)?,

                _ => (),
            }
        }

        // target_x and target_y depend on order.
        for (attr, value) in attrs.iter() {
            match attr.expanded() {
                expanded_name!("", "targetX") => {
                    self.target_x = {
                        let v = attr.parse_and_validate(value, |v: i32| {
                            if v >= 0 && v < self.order.0 as i32 {
                                Ok(v)
                            } else {
                                Err(ValueErrorKind::value_error(
                                    "targetX must be greater or equal to zero and less than orderX",
                                ))
                            }
                        })?;
                        Some(v as u32)
                    }
                }
                expanded_name!("", "targetY") => {
                    self.target_y = {
                        let v = attr.parse_and_validate(value, |v: i32| {
                            if v >= 0 && v < self.order.1 as i32 {
                                Ok(v)
                            } else {
                                Err(ValueErrorKind::value_error(
                                    "targetY must be greater or equal to zero and less than orderY",
                                ))
                            }
                        })?;
                        Some(v as u32)
                    }
                }
                _ => (),
            }
        }

        // Default values for target_x and target_y.
        if self.target_x.is_none() {
            self.target_x = Some(self.order.0 / 2);
        }
        if self.target_y.is_none() {
            self.target_y = Some(self.order.1 / 2);
        }

        // Finally, parse the kernel matrix.
        for (attr, value) in attrs
            .iter()
            .filter(|(attr, _)| attr.expanded() == expanded_name!("", "kernelMatrix"))
        {
            self.kernel_matrix = Some({
                let number_of_elements = self.order.0 as usize * self.order.1 as usize;

                // #352: Parse as an unbounded list rather than exact length to prevent aborts due
                //       to huge allocation attempts by underlying Vec::with_capacity().
                let NumberList(v) = NumberList::parse_str(value, NumberListLength::Unbounded)
                    .attribute(attr.clone())?;

                if v.len() != number_of_elements {
                    return Err(ValueErrorKind::value_error(&format!(
                        "incorrect number of elements: expected {}",
                        number_of_elements
                    )))
                    .attribute(attr);
                }

                DMatrix::from_data(VecStorage::new(
                    Dynamic::new(self.order.1 as usize),
                    Dynamic::new(self.order.0 as usize),
                    v,
                ))
            });
        }

        // kernel_matrix must have been specified.
        if self.kernel_matrix.is_none() {
            return Err(ValueErrorKind::value_error("the value must be set"))
                .attribute(QualName::new(None, ns!(svg), local_name!("kernelMatrix")));
        }

        // Default value for the divisor.
        if self.divisor.is_none() {
            self.divisor = Some(self.kernel_matrix.as_ref().unwrap().iter().sum());

            if self.divisor.unwrap() == 0.0 {
                self.divisor = Some(1.0);
            }
        }

        Ok(())
    }
}

impl FilterEffect for FeConvolveMatrix {
    fn render(
        &self,
        node: &Node,
        ctx: &FilterContext,
        acquired_nodes: &mut AcquiredNodes<'_>,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError> {
        #![allow(clippy::many_single_char_names)]

        let input = self.base.get_input(ctx, acquired_nodes, draw_ctx)?;
        let mut bounds = self
            .base
            .get_bounds(ctx, node.parent().as_ref())?
            .add_input(&input)
            .into_irect(draw_ctx);
        let original_bounds = bounds;

        let mut input_surface = if self.preserve_alpha {
            // preserve_alpha means we need to premultiply and unpremultiply the values.
            input.surface().unpremultiply(bounds)?
        } else {
            input.surface().clone()
        };

        let scale = self
            .kernel_unit_length
            .map(|(dx, dy)| ctx.paffine().transform_distance(dx, dy));

        if let Some((ox, oy)) = scale {
            // Scale the input surface to match kernel_unit_length.
            let (new_surface, new_bounds) = input_surface.scale(bounds, 1.0 / ox, 1.0 / oy)?;

            input_surface = new_surface;
            bounds = new_bounds;
        }

        let matrix = self.kernel_matrix.as_ref().unwrap();

        let mut surface = ExclusiveImageSurface::new(
            input_surface.width(),
            input_surface.height(),
            input.surface().surface_type(),
        )?;

        surface.modify(&mut |data, stride| {
            for (x, y, pixel) in Pixels::within(&input_surface, bounds) {
                // Compute the convolution rectangle bounds.
                let kernel_bounds = IRect::new(
                    x as i32 - self.target_x.unwrap() as i32,
                    y as i32 - self.target_y.unwrap() as i32,
                    x as i32 - self.target_x.unwrap() as i32 + self.order.0 as i32,
                    y as i32 - self.target_y.unwrap() as i32 + self.order.1 as i32,
                );

                // Do the convolution.
                let mut r = 0.0;
                let mut g = 0.0;
                let mut b = 0.0;
                let mut a = 0.0;

                for (x, y, pixel) in
                    PixelRectangle::within(&input_surface, bounds, kernel_bounds, self.edge_mode)
                {
                    let kernel_x = (kernel_bounds.x1 - x - 1) as usize;
                    let kernel_y = (kernel_bounds.y1 - y - 1) as usize;

                    r += f64::from(pixel.r) / 255.0 * matrix[(kernel_y, kernel_x)];
                    g += f64::from(pixel.g) / 255.0 * matrix[(kernel_y, kernel_x)];
                    b += f64::from(pixel.b) / 255.0 * matrix[(kernel_y, kernel_x)];

                    if !self.preserve_alpha {
                        a += f64::from(pixel.a) / 255.0 * matrix[(kernel_y, kernel_x)];
                    }
                }

                // If preserve_alpha is true, set a to the source alpha value.
                if self.preserve_alpha {
                    a = f64::from(pixel.a) / 255.0;
                } else {
                    a = a / self.divisor.unwrap() + self.bias;
                }

                let clamped_a = clamp(a, 0.0, 1.0);

                let compute = |x| {
                    let x = x / self.divisor.unwrap() + self.bias * a;

                    let x = if self.preserve_alpha {
                        // Premultiply the output value.
                        clamp(x, 0.0, 1.0) * clamped_a
                    } else {
                        clamp(x, 0.0, clamped_a)
                    };

                    ((x * 255.0) + 0.5) as u8
                };

                let output_pixel = Pixel {
                    r: compute(r),
                    g: compute(g),
                    b: compute(b),
                    a: ((clamped_a * 255.0) + 0.5) as u8,
                };

                data.set_pixel(stride, output_pixel, x, y);
            }
        });

        let mut surface = surface.share()?;

        if let Some((ox, oy)) = scale {
            // Scale the output surface back.
            surface = surface.scale_to(
                ctx.source_graphic().width(),
                ctx.source_graphic().height(),
                original_bounds,
                ox,
                oy,
            )?;

            bounds = original_bounds;
        }

        Ok(FilterResult {
            name: self.base.result.clone(),
            output: FilterOutput { surface, bounds },
        })
    }

    #[inline]
    fn is_affected_by_color_interpolation_filters(&self) -> bool {
        true
    }
}

impl Parse for EdgeMode {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<Self, ParseError<'i>> {
        Ok(parse_identifiers!(
            parser,
            "duplicate" => EdgeMode::Duplicate,
            "wrap" => EdgeMode::Wrap,
            "none" => EdgeMode::None,
        )?)
    }
}

// Used for the preserveAlpha attribute
impl Parse for bool {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<Self, ParseError<'i>> {
        Ok(parse_identifiers!(
            parser,
            "false" => false,
            "true" => true,
        )?)
    }
}
<|endoftext|>


<commit_before>extern crate clingo;

use std::env;
use clingo::*;

pub struct OnStatementData<'a> {
    atom: ClingoAstTerm,
    builder: &'a mut ClingoProgramBuilder,
}

struct MyAstHandler;
impl<'a> ClingoAstStatementHandler<OnStatementData<'a>> for MyAstHandler {
    // adds atom enable to all rule bodies
    fn on_statement(stm: &ClingoAstStatement, data: &mut OnStatementData) -> bool {

        // pass through all statements that are not rules
        if stm.get_type() != Ok(ClingoAstStatementType::Rule) {
            data.builder.add(stm).expect(
                "Failed to add statement to ProgramBuilder.",
            );
            return true;
        }

        // copy the current rule body
        let body = unsafe { stm.rule() }.body();
        let mut extended_body = std::vec::Vec::with_capacity(body.len() + 1);
        for e in body {
            extended_body.push(e.clone());
        }

        // create atom enable
        let lit = ClingoAstLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstLiteralType::Symbolic,
            &data.atom,
        );
        // add atom enable to the rule body
        let y: ClingoAstBodyLiteral = ClingoAstBodyLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstBodyLiteralType::Literal,
            &lit,
        );
        extended_body.push(y);

        // initialize the rule
        let head = unsafe { stm.rule() }.head();
        let rule = ClingoAstRule::new(head, &extended_body);

        // initialize the statement
        let stm2 = ClingoAstStatement::new_rule(stm.location(), &rule);

        // add the rewritten statement to the program
        data.builder.add(&stm2).expect(
            "Failed to add statement to ProgramBuilder.",
        );
        true
    }
}

fn print_model(model: &mut ClingoModel) {

    // retrieve the symbols in the model
    let atoms = model
        .symbols(ClingoShowType::Shown as clingo_show_type_bitset_t)
        .expect("Failed to retrieve symbols in the model.");

    print!("Model:");

    for atom in atoms {
        // retrieve and print the symbol's string
        print!(" {}", atom.to_string().unwrap());
    }
    println!();
}

fn solve(ctl: &mut ClingoControl) {

    let solve_mode = ClingoSolveMode::Yield;
    let assumptions = vec![];

    // get a solve handle
    let handle = ctl.solve(solve_mode, &assumptions).expect(
        "Failed retrieving solve handle.",
    );

    // loop over all models
    loop {
        handle.resume().expect("Failed resume on solve handle.");
        match handle.model() {
            // print the model
            Ok(model) => print_model(model),
            // stop if there are no more models
            Err(_) => break,
        }
    }

    // close the solve handle
    handle.get().expect(
        "Failed to get result from solve handle.",
    );
    handle.close().expect("Failed to close solve handle.");
}

fn get_modified_rules(builder: &mut ClingoProgramBuilder, sym: ClingoSymbol) {

        // begin building a program
        builder.begin().expect("Failed building logic program.");

        // initialize the location
        let location = create_clingo_location(0, 0, 0, 0, "<rewrite>", "<rewrite>");

        // initilize atom to add
        let atom = ClingoAstTerm::new_symbol(location, sym);

        let mut data = OnStatementData {
            atom: atom,
            builder: builder,
        };

        // get the AST of the program
        parse_program_with_event_handler("a :- not b. b :- not a.", &MyAstHandler, &mut data)
            .expect("Failed to parse logic program.");

        // add the external statement: #external enable.
        let ext = ClingoAstExternal::new(atom, &[]);

        let stm = ClingoAstStatement::new_external(location, ClingoAstStatementType::External, &ext);
        data.builder.add(&stm).expect(
            "Failed to add statement to ProgramBuilder.",
        );

        // finish building a program
        data.builder.end().expect(
            "Failed to finish building a program.",
        );
}

fn main() {

    // collect clingo options from the command line
    let options = env::args().skip(1).collect();

    let mut ctl = ClingoControl::new(options, 20).expect("Failed creating ClingoControl.");

    let sym = create_id("enable", true).unwrap();

    let sym2 = sym.clone();

    // get the program builder
    { let mut builder = ctl.program_builder().unwrap();
      get_modified_rules(&mut builder, sym);
    }
    // ground the base part
    let part = ClingoPart::new_part("base", &[]);
    let parts = vec![part];
    ctl.ground(parts).expect(
        "Failed to ground a logic program.",
    );

    // solve with external enable = false
    println!("Solving with enable = false...");
    solve(&mut ctl);

    // solve with external enable = true
    println!("Solving with enable = true...");
    ctl.assign_external(&sym2, ClingoTruthValue::True).expect(
        "Failed to assign #external enable true.",
    );
    solve(&mut ctl);

    // solve with external enable = false
    println!("Solving with enable = false...");
    ctl.assign_external(&sym2, ClingoTruthValue::False).expect(
        "Failed to assign #external enable false.",
    );
    solve(&mut ctl);
}
<commit_msg>remove function modify rules<commit_after>extern crate clingo;

use std::env;
use clingo::*;

pub struct OnStatementData<'a> {
    atom: ClingoAstTerm,
    builder: &'a mut ClingoProgramBuilder,
}

struct MyAstHandler;
impl<'a> ClingoAstStatementHandler<OnStatementData<'a>> for MyAstHandler {
    // adds atom enable to all rule bodies
    fn on_statement(stm: &ClingoAstStatement, data: &mut OnStatementData) -> bool {
        // pass through all statements that are not rules
        if stm.get_type() != Ok(ClingoAstStatementType::Rule) {
            data.builder
                .add(stm)
                .expect("Failed to add statement to ProgramBuilder.");
            return true;
        }

        // copy the current rule body
        let body = unsafe { stm.rule() }.body();
        let mut extended_body = std::vec::Vec::with_capacity(body.len() + 1);
        for e in body {
            extended_body.push(e.clone());
        }

        // create atom enable
        let lit = ClingoAstLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstLiteralType::Symbolic,
            &data.atom,
        );
        // add atom enable to the rule body
        let y: ClingoAstBodyLiteral = ClingoAstBodyLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstBodyLiteralType::Literal,
            &lit,
        );
        extended_body.push(y);

        // initialize the rule
        let head = unsafe { stm.rule() }.head();
        let rule = ClingoAstRule::new(head, &extended_body);

        // initialize the statement
        let stm2 = ClingoAstStatement::new_rule(stm.location(), &rule);

        // add the rewritten statement to the program
        data.builder
            .add(&stm2)
            .expect("Failed to add statement to ProgramBuilder.");
        true
    }
}

fn print_model(model: &mut ClingoModel) {
    // retrieve the symbols in the model
    let atoms = model
        .symbols(ClingoShowType::Shown as clingo_show_type_bitset_t)
        .expect("Failed to retrieve symbols in the model.");

    print!("Model:");

    for atom in atoms {
        // retrieve and print the symbol's string
        print!(" {}", atom.to_string().unwrap());
    }
    println!();
}

fn solve(ctl: &mut ClingoControl) {
    let solve_mode = ClingoSolveMode::Yield;
    let assumptions = vec![];

    // get a solve handle
    let handle = ctl.solve(solve_mode, &assumptions)
        .expect("Failed retrieving solve handle.");

    // loop over all models
    loop {
        handle.resume().expect("Failed resume on solve handle.");
        match handle.model() {
            // print the model
            Ok(model) => print_model(model),
            // stop if there are no more models
            Err(_) => break,
        }
    }

    // close the solve handle
    handle
        .get()
        .expect("Failed to get result from solve handle.");
    handle.close().expect("Failed to close solve handle.");
}

fn main() {
    // collect clingo options from the command line
    let options = env::args().skip(1).collect();

    let mut ctl = ClingoControl::new(options, 20).expect("Failed creating ClingoControl.");

    let sym = create_id("enable", true).unwrap();

    let sym2 = sym.clone();

    {
        // get the program builder
        let builder = ctl.program_builder().unwrap();

        // begin building a program
        builder.begin().expect("Failed building logic program.");

        // initialize the location
        let location = create_clingo_location(0, 0, 0, 0, "<rewrite>", "<rewrite>");

        // initilize atom to add
        let atom = ClingoAstTerm::new_symbol(location, sym);

        let mut data = OnStatementData {
            atom: atom,
            builder: builder,
        };

        // get the AST of the program
        parse_program_with_event_handler("a :- not b. b :- not a.", &MyAstHandler, &mut data)
            .expect("Failed to parse logic program.");

        // add the external statement: #external enable.
        let ext = ClingoAstExternal::new(atom, &[]);

        let stm =
            ClingoAstStatement::new_external(location, ClingoAstStatementType::External, &ext);
        data.builder
            .add(&stm)
            .expect("Failed to add statement to ProgramBuilder.");

        // finish building a program
        data.builder
            .end()
            .expect("Failed to finish building a program.");
    }
    // ground the base part
    let part = ClingoPart::new_part("base", &[]);
    let parts = vec![part];
    ctl.ground(parts)
        .expect("Failed to ground a logic program.");

    // solve with external enable = false
    println!("Solving with enable = false...");
    solve(&mut ctl);

    // solve with external enable = true
    println!("Solving with enable = true...");
    ctl.assign_external(&sym2, ClingoTruthValue::True)
        .expect("Failed to assign #external enable true.");
    solve(&mut ctl);

    // solve with external enable = false
    println!("Solving with enable = false...");
    ctl.assign_external(&sym2, ClingoTruthValue::False)
        .expect("Failed to assign #external enable false.");
    solve(&mut ctl);
}
<|endoftext|>