<commit_before>use num_bigint::BigInt;
use num_traits::ToPrimitive;
use std::cmp::max;
use std::str::FromStr;

mod precision;
use self::precision::*;

pub mod asset_symbol;
pub use self::asset_symbol::*;

pub mod cmp;
pub use self::cmp::*;

pub mod error;
pub use self::error::*;

pub const MAX_STR_LEN: usize = 32;
pub const MAX_PRECISION: u8 = 4;

pub const EMPTY_GOLD: Asset = Asset {
    amount: 0,
    decimals: 0,
    symbol: AssetSymbol::GOLD,
};

pub const EMPTY_SILVER: Asset = Asset {
    amount: 0,
    decimals: 0,
    symbol: AssetSymbol::SILVER,
};

#[derive(Clone, Debug)]
pub struct Asset {
    pub amount: i64,
    pub decimals: u8,
    pub symbol: AssetSymbol,
}

impl Asset {
    #[inline]
    pub fn checked_new(amount: i64, decimals: u8, symbol: AssetSymbol) -> Option<Asset> {
        if decimals > MAX_PRECISION {
            return None;
        }
        Some(Asset {
            amount,
            decimals,
            symbol,
        })
    }

    pub fn add(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol {
            return None;
        }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: a.checked_add(b)?,
            decimals,
            symbol: self.symbol,
        })
    }

    pub fn sub(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol {
            return None;
        }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: a.checked_sub(b)?,
            decimals,
            symbol: self.symbol,
        })
    }

    pub fn mul(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol || precision > MAX_PRECISION {
            return None;
        }
        let decimals = self.decimals + other.decimals;

        let mul = i128::from(self.amount).checked_mul(i128::from(other.amount))?;
        let final_mul = set_decimals_i128(mul, decimals, precision)?;
        if final_mul > i128::from(::std::i64::MAX) {
            return None;
        }
        Some(Asset {
            amount: final_mul as i64,
            decimals: precision,
            symbol: self.symbol,
        })
    }

    pub fn div(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol || other.amount == 0 || precision > MAX_PRECISION {
            return None;
        }
        let decimals = max(max(self.decimals, other.decimals), precision);
        let a = set_decimals_i64(self.amount, self.decimals, decimals * 2)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: set_decimals_i64(a.checked_div(b)?, decimals, precision)?,
            decimals: precision,
            symbol: self.symbol,
        })
    }

    pub fn pow(&self, num: u16, precision: u8) -> Option<Self> {
        if precision > MAX_PRECISION {
            return None;
        }
        if num == 0 {
            return Some(Asset {
                amount: set_decimals_i64(1, 0, precision)?,
                decimals: precision,
                symbol: self.symbol,
            });
        }
        let big_zero = BigInt::from(0);
        let big_one = BigInt::from(1);

        let decimals = u16::from(self.decimals).checked_mul(num)?;
        let mut res = BigInt::from(1);
        {
            let mut base = BigInt::from(self.amount);
            let mut exp = BigInt::from(num);
            loop {
                if &exp & &big_one == big_one {
                    res = &res * &base;
                }
                exp >>= 1;
                if exp == big_zero {
                    break;
                };
                base = &base * &base;
            }
        }

        res = set_decimals_big(&res, decimals, u16::from(precision));
        Some(Asset {
            amount: res.to_i64()?,
            decimals: precision,
            symbol: self.symbol,
        })
    }
}

impl ToString for Asset {
    fn to_string(&self) -> String {
        let mut s = self.amount.to_string();
        if self.decimals > 0 {
            let len = s.len();
            if len < self.decimals as usize {
                let start = if self.amount < 0 { 1 } else { 0 };
                let diff = self.decimals as usize - len + start;
                s.insert_str(start, "0.");
                s.insert_str(start + 2, &"0".repeat(diff));
            } else if len == self.decimals as usize {
                s.insert_str(0, "0.");
            } else {
                s.insert(len - (self.decimals as usize), '.');
            }
        }
        s.push(' ');
        s.push_str(self.symbol.as_str());
        s
    }
}

impl FromStr for Asset {
    type Err = AssetError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.len() > MAX_STR_LEN {
            return Err(AssetError {
                kind: AssetErrorKind::StrTooLarge,
            });
        }
        let mut split = s.trim().splitn(2, ' ');

        let amount: i64;
        let mut decimals: u8 = 0;
        match split.next() {
            Some(x) => {
                if let Some(pos) = x.find('.') {
                    let len = x.len() - 1;
                    if pos > 0 {
                        decimals = (len - pos) as u8;
                    } else {
                        decimals = len as u8;
                    }

                    if decimals > MAX_PRECISION {
                        return Err(AssetError {
                            kind: AssetErrorKind::InvalidAmount,
                        });
                    }

                    amount = match x.replace('.', "").parse() {
                        Ok(x) => x,
                        Err(_) => {
                            return Err(AssetError {
                                kind: AssetErrorKind::InvalidAmount,
                            });
                        }
                    }
                } else {
                    amount = match x.parse() {
                        Ok(x) => x,
                        Err(_) => {
                            return Err(AssetError {
                                kind: AssetErrorKind::InvalidAmount,
                            });
                        }
                    }
                }
            }
            None => {
                return Err(AssetError {
                    kind: AssetErrorKind::InvalidFormat,
                });
            }
        };

        let symbol = match split.next() {
            Some(x) => match AssetSymbol::parse_str(x) {
                Some(x) => x,
                None => {
                    return Err(AssetError {
                        kind: AssetErrorKind::InvalidAssetType,
                    });
                }
            },
            None => {
                return Err(AssetError {
                    kind: AssetErrorKind::InvalidFormat,
                });
            }
        };

        Ok(Asset {
            amount,
            decimals,
            symbol,
        })
    }
}

#[derive(Clone, Debug)]
pub struct Balance {
    pub gold: Asset,
    pub silver: Asset,
}

macro_rules! agnostic_op {
    ($op:ident) => {
        impl Balance {
            #[inline]
            #[must_use = "operation can fail"]
            pub fn $op(&mut self, asset: &Asset) -> Option<&mut Balance> {
                match asset.symbol {
                    AssetSymbol::GOLD => {
                        self.gold = self.gold.$op(asset)?;
                    }
                    AssetSymbol::SILVER => {
                        self.silver = self.silver.$op(asset)?;
                    }
                }
                Some(self)
            }
        }
    };
}

agnostic_op!(add);
agnostic_op!(sub);

impl Default for Balance {
    fn default() -> Balance {
        Balance {
            gold: EMPTY_GOLD,
            silver: EMPTY_SILVER,
        }
    }
}

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

    #[test]
    fn test_parse_valid_input() {
        let c = |asset: Asset, amount: &str, decimals: u8, symbol: AssetSymbol| {
            assert_eq!(asset.amount.to_string(), amount);
            assert_eq!(asset.decimals, decimals);
            assert_eq!(asset.symbol, symbol);
        };

        c(get_asset("1 GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset("1. GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset(".1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("-.1 GOLD"), "-1", 1, AssetSymbol::GOLD);
        c(get_asset("0.1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("1.0 SILVER"), "10", 1, AssetSymbol::SILVER);
        c(get_asset("0 SILVER"), "0", 0, AssetSymbol::SILVER);
        c(get_asset("-0.0 SILVER"), "0", 1, AssetSymbol::SILVER);
        c(get_asset("-1.0 SILVER"), "-10", 1, AssetSymbol::SILVER);
    }

    #[test]
    fn test_asset_to_str() {
        let c = |asset: Asset, s: &str| {
            assert_eq!(asset.to_string(), s);
        };
        c(get_asset("1.0001 GOLD"), "1.0001 GOLD");
        c(get_asset("0.0001 GOLD"), "0.0001 GOLD");
        c(get_asset("-0.0001 GOLD"), "-0.0001 GOLD");
        c(get_asset(".0001 GOLD"), "0.0001 GOLD");
        c(get_asset(".1 GOLD"), "0.1 GOLD");
        c(get_asset("1.0 GOLD"), "1.0 GOLD");
    }

    #[test]
    fn test_fail_parsing_invalid_input() {
        let c = |asset: &str, err: AssetErrorKind| {
            let e = Asset::from_str(asset).err().unwrap();
            assert_eq!(e.kind, err);
        };

        c("1e10 GOLD", AssetErrorKind::InvalidAmount);
        c("a100 GOLD", AssetErrorKind::InvalidAmount);
        c("100a GOLD", AssetErrorKind::InvalidAmount);

        c(
            "1234567890123456789012345678 GOLD",
            AssetErrorKind::StrTooLarge,
        );
        c("1", AssetErrorKind::InvalidFormat);

        c("1.0 GOLD a", AssetErrorKind::InvalidAssetType);
        c("1.0 gold", AssetErrorKind::InvalidAssetType);
    }

    #[test]
    fn test_set_precision() {
        let a = get_asset("1.5678 GOLD");
        assert_eq!(a.decimals, 4);
        assert_eq!(a.amount.to_string(), "15678");

        let a = a.mul(&get_asset("10000 GOLD"), 0).unwrap();
        assert_eq!(a.decimals, 0);
        assert_eq!(a.amount.to_string(), "15678");

        let a = a.div(&get_asset("100 GOLD"), 0).unwrap();
        assert_eq!(a.decimals, 0);
        assert_eq!(a.amount.to_string(), "156");

        let a = a.div(&get_asset("100 GOLD"), 2).unwrap();
        assert_eq!(a.decimals, 2);
        assert_eq!(a.amount.to_string(), "156");

        let a = a.div(&get_asset("100 GOLD"), 2).unwrap();
        assert_eq!(a.decimals, 2);
        assert_eq!(a.amount.to_string(), "1");
    }

    #[test]
    fn test_perform_arithmetic() {
        let c = |asset: &Asset, amount: &str| {
            assert_eq!(asset.to_string(), amount);
        };

        let a = get_asset("123.456 GOLD");
        c(&a.add(&get_asset("2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(&a.add(&get_asset("-2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.add(&get_asset(".0001 GOLD")).unwrap(), "123.4561 GOLD");
        c(&a.sub(&get_asset("2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.sub(&get_asset("-2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(
            &a.mul(&get_asset("100000.1111 GOLD"), 4).unwrap(),
            "12345613.7159 GOLD",
        );
        c(
            &a.mul(&get_asset("-100000.1111 GOLD"), 4).unwrap(),
            "-12345613.7159 GOLD",
        );
        c(&a.div(&get_asset("23 GOLD"), 3).unwrap(), "5.367 GOLD");
        c(&a.div(&get_asset("-23 GOLD"), 4).unwrap(), "-5.3676 GOLD");
        c(&a.pow(2, 4).unwrap(), "15241.3839 GOLD");
        c(&a.pow(3, 4).unwrap(), "1881640.2952 GOLD");
        c(&a, "123.456 GOLD");

        c(
            &get_asset("1.0002 GOLD").pow(1000, 4).unwrap(),
            "1.2213 GOLD",
        );
        c(
            &get_asset("10 GOLD").div(&get_asset("2 GOLD"), 0).unwrap(),
            "5 GOLD",
        );
        c(
            &get_asset("5 GOLD").div(&get_asset("10 GOLD"), 1).unwrap(),
            "0.5 GOLD",
        );

        assert!(&a.div(&get_asset("0 GOLD"), 1).is_none());
    }

    #[test]
    fn test_invalid_arithmetic() {
        let a = &get_asset("10 GOLD");
        let b = &get_asset("10 SILVER");

        assert!(a.add(b).is_none());
        assert!(a.sub(b).is_none());
        assert!(a.div(b, 8).is_none());
        assert!(a.mul(b, 8).is_none());
    }

    fn get_asset(s: &str) -> Asset {
        Asset::from_str(s).unwrap()
    }
}
<commit_msg>Optimize asset pow function slightly<commit_after>use num_bigint::BigInt;
use num_traits::ToPrimitive;
use std::cmp::max;
use std::str::FromStr;

mod precision;
use self::precision::*;

pub mod asset_symbol;
pub use self::asset_symbol::*;

pub mod cmp;
pub use self::cmp::*;

pub mod error;
pub use self::error::*;

pub const MAX_STR_LEN: usize = 32;
pub const MAX_PRECISION: u8 = 4;

pub const EMPTY_GOLD: Asset = Asset {
    amount: 0,
    decimals: 0,
    symbol: AssetSymbol::GOLD,
};

pub const EMPTY_SILVER: Asset = Asset {
    amount: 0,
    decimals: 0,
    symbol: AssetSymbol::SILVER,
};

#[derive(Clone, Debug)]
pub struct Asset {
    pub amount: i64,
    pub decimals: u8,
    pub symbol: AssetSymbol,
}

impl Asset {
    #[inline]
    pub fn checked_new(amount: i64, decimals: u8, symbol: AssetSymbol) -> Option<Asset> {
        if decimals > MAX_PRECISION {
            return None;
        }
        Some(Asset {
            amount,
            decimals,
            symbol,
        })
    }

    pub fn add(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol {
            return None;
        }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: a.checked_add(b)?,
            decimals,
            symbol: self.symbol,
        })
    }

    pub fn sub(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol {
            return None;
        }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: a.checked_sub(b)?,
            decimals,
            symbol: self.symbol,
        })
    }

    pub fn mul(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol || precision > MAX_PRECISION {
            return None;
        }
        let decimals = self.decimals + other.decimals;

        let mul = i128::from(self.amount).checked_mul(i128::from(other.amount))?;
        let final_mul = set_decimals_i128(mul, decimals, precision)?;
        if final_mul > i128::from(::std::i64::MAX) {
            return None;
        }
        Some(Asset {
            amount: final_mul as i64,
            decimals: precision,
            symbol: self.symbol,
        })
    }

    pub fn div(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol || other.amount == 0 || precision > MAX_PRECISION {
            return None;
        }
        let decimals = max(max(self.decimals, other.decimals), precision);
        let a = set_decimals_i64(self.amount, self.decimals, decimals * 2)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: set_decimals_i64(a.checked_div(b)?, decimals, precision)?,
            decimals: precision,
            symbol: self.symbol,
        })
    }

    pub fn pow(&self, num: u16, precision: u8) -> Option<Self> {
        if precision > MAX_PRECISION {
            return None;
        }
        if num == 0 {
            return Some(Asset {
                amount: set_decimals_i64(1, 0, precision)?,
                decimals: precision,
                symbol: self.symbol,
            });
        }

        let decimals = u16::from(self.decimals).checked_mul(num)?;
        let mut res = BigInt::from(1);
        {
            let mut base = BigInt::from(self.amount);
            let mut exp = num;
            loop {
                if exp & 1 == 1 {
                    res = &res * &base;
                }
                exp >>= 1;
                if exp == 0 {
                    break;
                };
                base = &base * &base;
            }
        }

        res = set_decimals_big(&res, decimals, u16::from(precision));
        Some(Asset {
            amount: res.to_i64()?,
            decimals: precision,
            symbol: self.symbol,
        })
    }
}

impl ToString for Asset {
    fn to_string(&self) -> String {
        let mut s = self.amount.to_string();
        if self.decimals > 0 {
            let len = s.len();
            if len < self.decimals as usize {
                let start = if self.amount < 0 { 1 } else { 0 };
                let diff = self.decimals as usize - len + start;
                s.insert_str(start, "0.");
                s.insert_str(start + 2, &"0".repeat(diff));
            } else if len == self.decimals as usize {
                s.insert_str(0, "0.");
            } else {
                s.insert(len - (self.decimals as usize), '.');
            }
        }
        s.push(' ');
        s.push_str(self.symbol.as_str());
        s
    }
}

impl FromStr for Asset {
    type Err = AssetError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.len() > MAX_STR_LEN {
            return Err(AssetError {
                kind: AssetErrorKind::StrTooLarge,
            });
        }
        let mut split = s.trim().splitn(2, ' ');

        let amount: i64;
        let mut decimals: u8 = 0;
        match split.next() {
            Some(x) => {
                if let Some(pos) = x.find('.') {
                    let len = x.len() - 1;
                    if pos > 0 {
                        decimals = (len - pos) as u8;
                    } else {
                        decimals = len as u8;
                    }

                    if decimals > MAX_PRECISION {
                        return Err(AssetError {
                            kind: AssetErrorKind::InvalidAmount,
                        });
                    }

                    amount = match x.replace('.', "").parse() {
                        Ok(x) => x,
                        Err(_) => {
                            return Err(AssetError {
                                kind: AssetErrorKind::InvalidAmount,
                            });
                        }
                    }
                } else {
                    amount = match x.parse() {
                        Ok(x) => x,
                        Err(_) => {
                            return Err(AssetError {
                                kind: AssetErrorKind::InvalidAmount,
                            });
                        }
                    }
                }
            }
            None => {
                return Err(AssetError {
                    kind: AssetErrorKind::InvalidFormat,
                });
            }
        };

        let symbol = match split.next() {
            Some(x) => match AssetSymbol::parse_str(x) {
                Some(x) => x,
                None => {
                    return Err(AssetError {
                        kind: AssetErrorKind::InvalidAssetType,
                    });
                }
            },
            None => {
                return Err(AssetError {
                    kind: AssetErrorKind::InvalidFormat,
                });
            }
        };

        Ok(Asset {
            amount,
            decimals,
            symbol,
        })
    }
}

#[derive(Clone, Debug)]
pub struct Balance {
    pub gold: Asset,
    pub silver: Asset,
}

macro_rules! agnostic_op {
    ($op:ident) => {
        impl Balance {
            #[inline]
            #[must_use = "operation can fail"]
            pub fn $op(&mut self, asset: &Asset) -> Option<&mut Balance> {
                match asset.symbol {
                    AssetSymbol::GOLD => {
                        self.gold = self.gold.$op(asset)?;
                    }
                    AssetSymbol::SILVER => {
                        self.silver = self.silver.$op(asset)?;
                    }
                }
                Some(self)
            }
        }
    };
}

agnostic_op!(add);
agnostic_op!(sub);

impl Default for Balance {
    fn default() -> Balance {
        Balance {
            gold: EMPTY_GOLD,
            silver: EMPTY_SILVER,
        }
    }
}

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

    #[test]
    fn test_parse_valid_input() {
        let c = |asset: Asset, amount: &str, decimals: u8, symbol: AssetSymbol| {
            assert_eq!(asset.amount.to_string(), amount);
            assert_eq!(asset.decimals, decimals);
            assert_eq!(asset.symbol, symbol);
        };

        c(get_asset("1 GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset("1. GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset(".1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("-.1 GOLD"), "-1", 1, AssetSymbol::GOLD);
        c(get_asset("0.1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("1.0 SILVER"), "10", 1, AssetSymbol::SILVER);
        c(get_asset("0 SILVER"), "0", 0, AssetSymbol::SILVER);
        c(get_asset("-0.0 SILVER"), "0", 1, AssetSymbol::SILVER);
        c(get_asset("-1.0 SILVER"), "-10", 1, AssetSymbol::SILVER);
    }

    #[test]
    fn test_asset_to_str() {
        let c = |asset: Asset, s: &str| {
            assert_eq!(asset.to_string(), s);
        };
        c(get_asset("1.0001 GOLD"), "1.0001 GOLD");
        c(get_asset("0.0001 GOLD"), "0.0001 GOLD");
        c(get_asset("-0.0001 GOLD"), "-0.0001 GOLD");
        c(get_asset(".0001 GOLD"), "0.0001 GOLD");
        c(get_asset(".1 GOLD"), "0.1 GOLD");
        c(get_asset("1.0 GOLD"), "1.0 GOLD");
    }

    #[test]
    fn test_fail_parsing_invalid_input() {
        let c = |asset: &str, err: AssetErrorKind| {
            let e = Asset::from_str(asset).err().unwrap();
            assert_eq!(e.kind, err);
        };

        c("1e10 GOLD", AssetErrorKind::InvalidAmount);
        c("a100 GOLD", AssetErrorKind::InvalidAmount);
        c("100a GOLD", AssetErrorKind::InvalidAmount);

        c(
            "1234567890123456789012345678 GOLD",
            AssetErrorKind::StrTooLarge,
        );
        c("1", AssetErrorKind::InvalidFormat);

        c("1.0 GOLD a", AssetErrorKind::InvalidAssetType);
        c("1.0 gold", AssetErrorKind::InvalidAssetType);
    }

    #[test]
    fn test_set_precision() {
        let a = get_asset("1.5678 GOLD");
        assert_eq!(a.decimals, 4);
        assert_eq!(a.amount.to_string(), "15678");

        let a = a.mul(&get_asset("10000 GOLD"), 0).unwrap();
        assert_eq!(a.decimals, 0);
        assert_eq!(a.amount.to_string(), "15678");

        let a = a.div(&get_asset("100 GOLD"), 0).unwrap();
        assert_eq!(a.decimals, 0);
        assert_eq!(a.amount.to_string(), "156");

        let a = a.div(&get_asset("100 GOLD"), 2).unwrap();
        assert_eq!(a.decimals, 2);
        assert_eq!(a.amount.to_string(), "156");

        let a = a.div(&get_asset("100 GOLD"), 2).unwrap();
        assert_eq!(a.decimals, 2);
        assert_eq!(a.amount.to_string(), "1");
    }

    #[test]
    fn test_perform_arithmetic() {
        let c = |asset: &Asset, amount: &str| {
            assert_eq!(asset.to_string(), amount);
        };

        let a = get_asset("123.456 GOLD");
        c(&a.add(&get_asset("2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(&a.add(&get_asset("-2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.add(&get_asset(".0001 GOLD")).unwrap(), "123.4561 GOLD");
        c(&a.sub(&get_asset("2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.sub(&get_asset("-2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(
            &a.mul(&get_asset("100000.1111 GOLD"), 4).unwrap(),
            "12345613.7159 GOLD",
        );
        c(
            &a.mul(&get_asset("-100000.1111 GOLD"), 4).unwrap(),
            "-12345613.7159 GOLD",
        );
        c(&a.div(&get_asset("23 GOLD"), 3).unwrap(), "5.367 GOLD");
        c(&a.div(&get_asset("-23 GOLD"), 4).unwrap(), "-5.3676 GOLD");
        c(&a.pow(2, 4).unwrap(), "15241.3839 GOLD");
        c(&a.pow(3, 4).unwrap(), "1881640.2952 GOLD");
        c(&a, "123.456 GOLD");

        c(
            &get_asset("1.0002 GOLD").pow(1000, 4).unwrap(),
            "1.2213 GOLD",
        );
        c(
            &get_asset("10 GOLD").div(&get_asset("2 GOLD"), 0).unwrap(),
            "5 GOLD",
        );
        c(
            &get_asset("5 GOLD").div(&get_asset("10 GOLD"), 1).unwrap(),
            "0.5 GOLD",
        );

        assert!(&a.div(&get_asset("0 GOLD"), 1).is_none());
    }

    #[test]
    fn test_invalid_arithmetic() {
        let a = &get_asset("10 GOLD");
        let b = &get_asset("10 SILVER");

        assert!(a.add(b).is_none());
        assert!(a.sub(b).is_none());
        assert!(a.div(b, 8).is_none());
        assert!(a.mul(b, 8).is_none());
    }

    fn get_asset(s: &str) -> Asset {
        Asset::from_str(s).unwrap()
    }
}
<|endoftext|>


<commit_before>//! See `AssistContext`

use algo::find_covering_element;
use hir::Semantics;
use ra_db::{FileId, FileRange};
use ra_fmt::{leading_indent, reindent};
use ra_ide_db::{
    source_change::{SourceChange, SourceFileEdit},
    RootDatabase,
};
use ra_syntax::{
    algo::{self, find_node_at_offset, SyntaxRewriter},
    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextSize,
    TokenAtOffset,
};
use ra_text_edit::TextEditBuilder;

use crate::{
    assist_config::{AssistConfig, SnippetCap},
    Assist, AssistId, GroupLabel, ResolvedAssist,
};
use rustc_hash::FxHashMap;

/// `AssistContext` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat over-engineered approach, given the current needs.
/// The assists workflow consists of two phases. In the first phase, a user asks
/// for the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then necessary to answer "is
///   assist applicable" in the first phase.
/// * second, when we are applying assist, we don't have a guarantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach reminiscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `resolve`. It first check if an edit is applicable (potentially computing
/// info required to compute the actual edit). If it is applicable, and
/// `resolve` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `resolve = false`, and then applying the selected edit again, with
/// `resolve = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagerly :-)
pub(crate) struct AssistContext<'a> {
    pub(crate) config: &'a AssistConfig,
    pub(crate) sema: Semantics<'a, RootDatabase>,
    pub(crate) db: &'a RootDatabase,
    pub(crate) frange: FileRange,
    source_file: SourceFile,
}

impl<'a> AssistContext<'a> {
    pub(crate) fn new(
        sema: Semantics<'a, RootDatabase>,
        config: &'a AssistConfig,
        frange: FileRange,
    ) -> AssistContext<'a> {
        let source_file = sema.parse(frange.file_id);
        let db = sema.db;
        AssistContext { config, sema, db, frange, source_file }
    }

    // NB, this ignores active selection.
    pub(crate) fn offset(&self) -> TextSize {
        self.frange.range.start()
    }

    pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
        self.source_file.syntax().token_at_offset(self.offset())
    }
    pub(crate) fn find_token_at_offset(&self, kind: SyntaxKind) -> Option<SyntaxToken> {
        self.token_at_offset().find(|it| it.kind() == kind)
    }
    pub(crate) fn find_node_at_offset<N: AstNode>(&self) -> Option<N> {
        find_node_at_offset(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn find_node_at_offset_with_descend<N: AstNode>(&self) -> Option<N> {
        self.sema.find_node_at_offset_with_descend(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn covering_element(&self) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), self.frange.range)
    }
    // FIXME: remove
    pub(crate) fn covering_node_for_range(&self, range: TextRange) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), range)
    }
}

pub(crate) struct Assists {
    resolve: bool,
    file: FileId,
    buf: Vec<(Assist, Option<SourceChange>)>,
}

impl Assists {
    pub(crate) fn new_resolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: true, file: ctx.frange.file_id, buf: Vec::new() }
    }
    pub(crate) fn new_unresolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: false, file: ctx.frange.file_id, buf: Vec::new() }
    }

    pub(crate) fn finish_unresolved(self) -> Vec<Assist> {
        assert!(!self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| {
                assert!(edit.is_none());
                label
            })
            .collect()
    }

    pub(crate) fn finish_resolved(self) -> Vec<ResolvedAssist> {
        assert!(self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| ResolvedAssist { assist: label, source_change: edit.unwrap() })
            .collect()
    }

    pub(crate) fn add(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl(label, f)
    }
    pub(crate) fn add_in_multiple_files(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl_multiple_files(label, f)
    }
    pub(crate) fn add_group(
        &mut self,
        group: &GroupLabel,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), Some(group.clone()), target);
        self.add_impl(label, f)
    }
    fn add_impl(&mut self, label: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
        let source_change = if self.resolve {
            let mut builder = AssistBuilder::new(self.file);
            f(&mut builder);
            Some(builder.finish())
        } else {
            None
        };

        self.buf.push((label, source_change));
        Some(())
    }

    fn add_impl_multiple_files(
        &mut self,
        label: Assist,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        if !self.resolve {
            return None;
        }
        let mut director = AssistDirector::default();
        f(&mut director);
        let changes = director.finish();
        let file_edits: Vec<SourceFileEdit> =
            changes.into_iter().map(|mut change| change.source_file_edits.pop().unwrap()).collect();

        let source_change = SourceChange {
            source_file_edits: file_edits,
            file_system_edits: vec![],
            is_snippet: false,
        };

        self.buf.push((label, Some(source_change)));
        Some(())
    }

    fn finish(mut self) -> Vec<(Assist, Option<SourceChange>)> {
        self.buf.sort_by_key(|(label, _edit)| label.target.len());
        self.buf
    }
}

pub(crate) struct AssistBuilder {
    edit: TextEditBuilder,
    file: FileId,
    is_snippet: bool,
}

impl AssistBuilder {
    pub(crate) fn new(file: FileId) -> AssistBuilder {
        AssistBuilder { edit: TextEditBuilder::default(), file, is_snippet: false }
    }

    /// Remove specified `range` of text.
    pub(crate) fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }
    /// Append specified `text` at the given `offset`
    pub(crate) fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }
    /// Append specified `snippet` at the given `offset`
    pub(crate) fn insert_snippet(
        &mut self,
        _cap: SnippetCap,
        offset: TextSize,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.insert(offset, snippet);
    }
    /// Replaces specified `range` of text with a given string.
    pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }
    /// Replaces specified `range` of text with a given `snippet`.
    pub(crate) fn replace_snippet(
        &mut self,
        _cap: SnippetCap,
        range: TextRange,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.replace(range, snippet);
    }
    pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
        algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
    }
    /// Replaces specified `node` of text with a given string, reindenting the
    /// string to maintain `node`'s existing indent.
    // FIXME: remove in favor of ra_syntax::edit::IndentLevel::increase_indent
    pub(crate) fn replace_node_and_indent(
        &mut self,
        node: &SyntaxNode,
        replace_with: impl Into<String>,
    ) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, &indent)
        }
        self.replace(node.text_range(), replace_with)
    }
    pub(crate) fn rewrite(&mut self, rewriter: SyntaxRewriter) {
        let node = rewriter.rewrite_root().unwrap();
        let new = rewriter.rewrite(&node);
        algo::diff(&node, &new).into_text_edit(&mut self.edit)
    }

    // FIXME: better API
    pub(crate) fn set_file(&mut self, assist_file: FileId) {
        self.file = assist_file;
    }

    // FIXME: kill this API
    /// Get access to the raw `TextEditBuilder`.
    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
        &mut self.edit
    }

    fn finish(self) -> SourceChange {
        let edit = self.edit.finish();
        let source_file_edit = SourceFileEdit { file_id: self.file, edit };
        let mut res: SourceChange = source_file_edit.into();
        if self.is_snippet {
            res.is_snippet = true;
        }
        res
    }
}

pub(crate) struct AssistDirector {
    builders: FxHashMap<FileId, AssistBuilder>,
}

impl AssistDirector {
    pub(crate) fn perform(&mut self, file_id: FileId, f: impl FnOnce(&mut AssistBuilder)) {
        let mut builder = self.builders.entry(file_id).or_insert(AssistBuilder::new(file_id));
        f(&mut builder);
    }

    fn finish(self) -> Vec<SourceChange> {
        self.builders
            .into_iter()
            .map(|(_, builder)| builder.finish())
            .collect::<Vec<SourceChange>>()
    }
}

impl Default for AssistDirector {
    fn default() -> Self {
        AssistDirector { builders: FxHashMap::default() }
    }
}
<commit_msg>Fix incorrect behaviour if not resolved<commit_after>//! See `AssistContext`

use algo::find_covering_element;
use hir::Semantics;
use ra_db::{FileId, FileRange};
use ra_fmt::{leading_indent, reindent};
use ra_ide_db::{
    source_change::{SourceChange, SourceFileEdit},
    RootDatabase,
};
use ra_syntax::{
    algo::{self, find_node_at_offset, SyntaxRewriter},
    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextSize,
    TokenAtOffset,
};
use ra_text_edit::TextEditBuilder;

use crate::{
    assist_config::{AssistConfig, SnippetCap},
    Assist, AssistId, GroupLabel, ResolvedAssist,
};
use rustc_hash::FxHashMap;

/// `AssistContext` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat over-engineered approach, given the current needs.
/// The assists workflow consists of two phases. In the first phase, a user asks
/// for the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then necessary to answer "is
///   assist applicable" in the first phase.
/// * second, when we are applying assist, we don't have a guarantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach reminiscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `resolve`. It first check if an edit is applicable (potentially computing
/// info required to compute the actual edit). If it is applicable, and
/// `resolve` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `resolve = false`, and then applying the selected edit again, with
/// `resolve = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagerly :-)
pub(crate) struct AssistContext<'a> {
    pub(crate) config: &'a AssistConfig,
    pub(crate) sema: Semantics<'a, RootDatabase>,
    pub(crate) db: &'a RootDatabase,
    pub(crate) frange: FileRange,
    source_file: SourceFile,
}

impl<'a> AssistContext<'a> {
    pub(crate) fn new(
        sema: Semantics<'a, RootDatabase>,
        config: &'a AssistConfig,
        frange: FileRange,
    ) -> AssistContext<'a> {
        let source_file = sema.parse(frange.file_id);
        let db = sema.db;
        AssistContext { config, sema, db, frange, source_file }
    }

    // NB, this ignores active selection.
    pub(crate) fn offset(&self) -> TextSize {
        self.frange.range.start()
    }

    pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
        self.source_file.syntax().token_at_offset(self.offset())
    }
    pub(crate) fn find_token_at_offset(&self, kind: SyntaxKind) -> Option<SyntaxToken> {
        self.token_at_offset().find(|it| it.kind() == kind)
    }
    pub(crate) fn find_node_at_offset<N: AstNode>(&self) -> Option<N> {
        find_node_at_offset(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn find_node_at_offset_with_descend<N: AstNode>(&self) -> Option<N> {
        self.sema.find_node_at_offset_with_descend(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn covering_element(&self) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), self.frange.range)
    }
    // FIXME: remove
    pub(crate) fn covering_node_for_range(&self, range: TextRange) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), range)
    }
}

pub(crate) struct Assists {
    resolve: bool,
    file: FileId,
    buf: Vec<(Assist, Option<SourceChange>)>,
}

impl Assists {
    pub(crate) fn new_resolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: true, file: ctx.frange.file_id, buf: Vec::new() }
    }
    pub(crate) fn new_unresolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: false, file: ctx.frange.file_id, buf: Vec::new() }
    }

    pub(crate) fn finish_unresolved(self) -> Vec<Assist> {
        assert!(!self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| {
                assert!(edit.is_none());
                label
            })
            .collect()
    }

    pub(crate) fn finish_resolved(self) -> Vec<ResolvedAssist> {
        assert!(self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| ResolvedAssist { assist: label, source_change: edit.unwrap() })
            .collect()
    }

    pub(crate) fn add(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl(label, f)
    }
    pub(crate) fn add_in_multiple_files(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl_multiple_files(label, f)
    }
    pub(crate) fn add_group(
        &mut self,
        group: &GroupLabel,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), Some(group.clone()), target);
        self.add_impl(label, f)
    }
    fn add_impl(&mut self, label: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
        let source_change = if self.resolve {
            let mut builder = AssistBuilder::new(self.file);
            f(&mut builder);
            Some(builder.finish())
        } else {
            None
        };

        self.buf.push((label, source_change));
        Some(())
    }

    fn add_impl_multiple_files(
        &mut self,
        label: Assist,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        if !self.resolve {
            self.buf.push((label, None));
            return None;
        }
        let mut director = AssistDirector::default();
        f(&mut director);
        let changes = director.finish();
        let file_edits: Vec<SourceFileEdit> =
            changes.into_iter().map(|mut change| change.source_file_edits.pop().unwrap()).collect();

        let source_change = SourceChange {
            source_file_edits: file_edits,
            file_system_edits: vec![],
            is_snippet: false,
        };

        self.buf.push((label, Some(source_change)));
        Some(())
    }

    fn finish(mut self) -> Vec<(Assist, Option<SourceChange>)> {
        self.buf.sort_by_key(|(label, _edit)| label.target.len());
        self.buf
    }
}

pub(crate) struct AssistBuilder {
    edit: TextEditBuilder,
    file: FileId,
    is_snippet: bool,
}

impl AssistBuilder {
    pub(crate) fn new(file: FileId) -> AssistBuilder {
        AssistBuilder { edit: TextEditBuilder::default(), file, is_snippet: false }
    }

    /// Remove specified `range` of text.
    pub(crate) fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }
    /// Append specified `text` at the given `offset`
    pub(crate) fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }
    /// Append specified `snippet` at the given `offset`
    pub(crate) fn insert_snippet(
        &mut self,
        _cap: SnippetCap,
        offset: TextSize,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.insert(offset, snippet);
    }
    /// Replaces specified `range` of text with a given string.
    pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }
    /// Replaces specified `range` of text with a given `snippet`.
    pub(crate) fn replace_snippet(
        &mut self,
        _cap: SnippetCap,
        range: TextRange,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.replace(range, snippet);
    }
    pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
        algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
    }
    /// Replaces specified `node` of text with a given string, reindenting the
    /// string to maintain `node`'s existing indent.
    // FIXME: remove in favor of ra_syntax::edit::IndentLevel::increase_indent
    pub(crate) fn replace_node_and_indent(
        &mut self,
        node: &SyntaxNode,
        replace_with: impl Into<String>,
    ) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, &indent)
        }
        self.replace(node.text_range(), replace_with)
    }
    pub(crate) fn rewrite(&mut self, rewriter: SyntaxRewriter) {
        let node = rewriter.rewrite_root().unwrap();
        let new = rewriter.rewrite(&node);
        algo::diff(&node, &new).into_text_edit(&mut self.edit)
    }

    // FIXME: better API
    pub(crate) fn set_file(&mut self, assist_file: FileId) {
        self.file = assist_file;
    }

    // FIXME: kill this API
    /// Get access to the raw `TextEditBuilder`.
    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
        &mut self.edit
    }

    fn finish(self) -> SourceChange {
        let edit = self.edit.finish();
        let source_file_edit = SourceFileEdit { file_id: self.file, edit };
        let mut res: SourceChange = source_file_edit.into();
        if self.is_snippet {
            res.is_snippet = true;
        }
        res
    }
}

pub(crate) struct AssistDirector {
    builders: FxHashMap<FileId, AssistBuilder>,
}

impl AssistDirector {
    pub(crate) fn perform(&mut self, file_id: FileId, f: impl FnOnce(&mut AssistBuilder)) {
        let mut builder = self.builders.entry(file_id).or_insert(AssistBuilder::new(file_id));
        f(&mut builder);
    }

    fn finish(self) -> Vec<SourceChange> {
        self.builders
            .into_iter()
            .map(|(_, builder)| builder.finish())
            .collect::<Vec<SourceChange>>()
    }
}

impl Default for AssistDirector {
    fn default() -> Self {
        AssistDirector { builders: FxHashMap::default() }
    }
}
<|endoftext|>


<commit_before>// vim: tw=80
/// Generate Expectation and Expectations types with specific signatures
///
/// # Arguments
///
/// * `expectation`:    Name of the generated Expectation type
/// * `module`:         Name of the private module to create
/// * `o`:              Output type.  Must be static
/// * `methty`:         Comma-delimited sequence of arguments types for the
///                     method being mocked.
/// * `matchcall`:      Comma-delimited sequence of expressions that produce
///                     values of type `matchty` from values of type `methty`.
/// * `args`:           comma-delimited sequence of argument names for each
///                     argument.  Ideally this wouldn't need to be specified,
///                     but it is due to Rust's macro hygiene rules.
/// * `predargs`:       Comma-delimited sequence of identifiers of the same
///                     length as `args`, but distinct.
/// * `matchty`:        comma-delimited sequence of types for each match
///                     argument.  Must all be `'static`.
///
/// # Examples
///
/// ```
/// // Mock a method like foo(&self, x: u32, y: &i16) -> u32
/// omnimock!{(u32, &i16), u32, i, [&i.0, i.1], [i0, i1], [u32, i16]}
#[macro_export]
macro_rules! omnimock {(
        $expectation:ident,
        $module:ident,
        $o:ty,
        [ $( $methty:ty ),* ],
        [ $( $matchcall:expr ),* ],
        [ $( $args:ident ),* ],
        [ $( $predargs:ident ),* ],
        [ $( $matchty:ty ),* ]) =>
    {
        mod $module {
            use super::*;

        enum Rfunc {
            Default,
            // Indicates that a `return_once` expectation has already returned
            Expired,
            Mut(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
            // Should be Box<dyn FnOnce> once that feature is stabilized
            // https://github.com/rust-lang/rust/issues/28796
            Once(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
        }

        impl  Rfunc {
            fn call_mut(&mut self, $( $args: $methty, )* ) -> $o {
                match self {
                    Rfunc::Default => {
                        unimplemented!()
                        //Self::return_default()
                    },
                    Rfunc::Expired => {
                        panic!("Called a method twice that was expected only once")
                    },
                    Rfunc::Mut(f) => {
                        f( $( $args, )* )
                    },
                    Rfunc::Once(_) => {
                        let fo = ::std::mem::replace(self, Rfunc::Expired);
                        if let Rfunc::Once(mut f) = fo {
                            f( $( $args, )* )
                        } else {
                            unreachable!()
                        }
                    },
                }
            }
        }

        impl std::default::Default for Rfunc {
            fn default() -> Self {
                Rfunc::Default
            }
        }

        enum Matcher {
            Func(Box<Fn($( &$matchty, )* ) -> bool>),
            Pred( $( Box<::mockall::Predicate<$matchty>>, )* )
        }

        impl Matcher {
            fn eval(&self, $( $args: &$matchty, )*) -> bool {
                match self {
                    Matcher::Func(f) => f($( $args, )*),
                    Matcher::Pred($( $predargs, )*) =>
                        [$( $predargs.eval($args), )*]
                        .into_iter()
                        .all(|x| *x),
                }
            }

            fn verify(&self, $( $args: &$matchty, )* ) {
                use predicates_tree::CaseTreeExt;
                match self {
                    Matcher::Func(f) => assert!(f($( $args, )*),
                        "Expectation didn't match arguments"),
                    Matcher::Pred($( $predargs, )*) => {
                        $(if let Some(c) = $predargs.find_case(false, $args) {
                            panic!("Expectation didn't match arguments:\n{}",
                                   c.tree());
                        })*
                    }
                }
            }
        }

        impl Default for Matcher {
            #[allow(unused_variables)]
            fn default() -> Self {
                Matcher::Func(Box::new(|$( $args, )*| true))
            }
        }

        #[derive(Default)]
        struct Common {
            matcher: ::std::sync::Mutex<Matcher>,
            seq_handle: Option<::mockall::SeqHandle>,
            times: ::mockall::Times
        }

        impl Common {
            fn call(&self, $( $args: &$matchty, )* ) {
                self.matcher.lock().unwrap().verify($( $args, )*);
                self.times.call();
                self.verify_sequence();
                if self.times.is_satisfied() {
                    self.satisfy_sequence()
                }
            }

            fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                assert!(self.times.is_exact(),
                    "Only Expectations with an exact call count have sequences");
                self.seq_handle = Some(seq.next());
                self
            }

            /// Forbid this expectation from ever being called
            fn never(&mut self) {
                self.times.never();
            }

            fn satisfy_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.satisfy()
                }
            }

            /// Require this expectation to be called exactly `n` times.
            fn times(&mut self, n: usize) {
                self.times.n(n);
            }

            /// Allow this expectation to be called any number of times
            fn times_any(&mut self) {
                self.times.any();
            }

            /// Allow this expectation to be called any number of times within a given
            /// range
            fn times_range(&mut self, range: ::std::ops::Range<usize>) {
                self.times.range(range);
            }

            fn verify_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.verify()
                }
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*)
            {
                use ::std::ops::DerefMut;
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Pred($( Box::new($args), )*);
                ::std::mem::replace(guard.deref_mut(), m);
            }

            fn withf<F>(&mut self, f: F)
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                use ::std::ops::DerefMut;
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Func(Box::new(f));
                ::std::mem::replace(guard.deref_mut(), m);
            }
        }

        #[derive(Default)]
        pub struct Expectation {
            common: Common,
            rfunc: ::std::sync::Mutex<Rfunc>,
        }

        impl Expectation {
            pub fn call(&self, $( $args: $methty, )* ) -> $o {
                self.common.call($( $matchcall, )*);
                self.rfunc.lock().unwrap().call_mut($( $args, )*)
            }

            pub fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                self.common.in_sequence(seq);
                self
            }

            /// Forbid this expectation from ever being called
            pub fn never(&mut self) -> &mut Self {
                self.common.never();
                self
            }

            /// Supply an `FnOnce` closure that will provide the return value
            /// for this Expectation.  This is useful for return types that
            /// aren't `Clone`.  It will be an error to call this Expectation
            /// multiple times.
            pub fn return_once<F>(&mut self, f: F) -> &mut Self
                where F: FnOnce($( $methty, )*) -> $o + Send + 'static
            {
                use ::std::ops::DerefMut;
                let mut fopt = Some(f);
                let fmut = move |$( $args: $methty, )*| {
                    if let Some(f) = fopt.take() {
                        f($( $args, )*)
                    } else {
                        panic!("Called a method twice that was expected only once")
                    }
                };
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(),
                                        Rfunc::Once(Box::new(fmut)));
                }
                self
            }

            /// Single-threaded version of [`return_once`](#method.return_once).
            /// This is useful for return types that are neither `Send` nor
            /// `Clone`.
            ///
            /// It is a runtime error to call the mock method from a different
            /// thread than the one that originally called this method.  It is
            /// also a runtime error to call the method more than once.
            pub fn return_once_st<F>(&mut self, f: F) -> &mut Self
                where F: FnOnce($( $methty, )*) -> $o + 'static
            {
                use ::std::ops::DerefMut;
                let mut fragile = Some(::fragile::Fragile::new(f));
                let fmut = Box::new(move |$( $args: $methty, )*| {
                    match fragile.take() {
                        Some(frag) => (frag.into_inner())($( $args, )*),
                        None => panic!(
                            "Called a method twice that was expected only once")
                    }
                });
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(), Rfunc::Once(fmut));
                }
                self
            }

            pub fn returning<F>(&mut self, f: F) -> &mut Self
                where F: FnMut($( $methty, )*) -> $o + Send + 'static
            {
                {
                    use ::std::ops::DerefMut;
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(), Rfunc::Mut(Box::new(f)));
                }
                self
            }

            /// Single-threaded version of [`returning`](#method.returning).
            /// Can be used when the argument or return type isn't `Send`.
            ///
            /// It is a runtime error to call the mock method from a different
            /// thread than the one that originally called this method.
            pub fn returning_st<F>(&mut self, f: F) -> &mut Self
                where F: FnMut($( $methty, )*) -> $o + 'static
            {
                use ::std::ops::DerefMut;
                let mut fragile = ::fragile::Fragile::new(f);
                let fmut = move |$( $args: $methty, )*| {
                    (fragile.get_mut())($( $args, )*)
                };
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(),
                                        Rfunc::Mut(Box::new(fmut)));
                }
                self
            }

            /// Require this expectation to be called exactly `n` times.
            pub fn times(&mut self, n: usize) -> &mut Self {
                self.common.times(n);
                self
            }

            /// Allow this expectation to be called any number of times
            ///
            /// This behavior is the default, but the method is provided in case the
            /// default behavior changes.
            pub fn times_any(&mut self) -> &mut Self {
                self.common.times_any();
                self
            }

            /// Allow this expectation to be called any number of times within a
            /// given range
            pub fn times_range(&mut self, range: ::std::ops::Range<usize>)
                -> &mut Self
            {
                self.common.times_range(range);
                self
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            pub fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*) -> &mut Self
            {
                self.common.with($( $args, )*);
                self
            }

            pub fn withf<F>(&mut self, f: F) -> &mut Self
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                self.common.withf(f);
                self
            }
        }

        }
        use $module::Expectation as $expectation;
    }
}
<commit_msg>omnimock: consolidate use statements<commit_after>// vim: tw=80
/// Generate Expectation and Expectations types with specific signatures
///
/// # Arguments
///
/// * `expectation`:    Name of the generated Expectation type
/// * `module`:         Name of the private module to create
/// * `o`:              Output type.  Must be static
/// * `methty`:         Comma-delimited sequence of arguments types for the
///                     method being mocked.
/// * `matchcall`:      Comma-delimited sequence of expressions that produce
///                     values of type `matchty` from values of type `methty`.
/// * `args`:           comma-delimited sequence of argument names for each
///                     argument.  Ideally this wouldn't need to be specified,
///                     but it is due to Rust's macro hygiene rules.
/// * `predargs`:       Comma-delimited sequence of identifiers of the same
///                     length as `args`, but distinct.
/// * `matchty`:        comma-delimited sequence of types for each match
///                     argument.  Must all be `'static`.
///
/// # Examples
///
/// ```
/// // Mock a method like foo(&self, x: u32, y: &i16) -> u32
/// omnimock!{(u32, &i16), u32, i, [&i.0, i.1], [i0, i1], [u32, i16]}
#[macro_export]
macro_rules! omnimock {(
        $expectation:ident,
        $module:ident,
        $o:ty,
        [ $( $methty:ty ),* ],
        [ $( $matchcall:expr ),* ],
        [ $( $args:ident ),* ],
        [ $( $predargs:ident ),* ],
        [ $( $matchty:ty ),* ]) =>
    {
        mod $module {
        use ::predicates_tree::CaseTreeExt;
        use ::std::ops::DerefMut;
        use super::*;

        enum Rfunc {
            Default,
            // Indicates that a `return_once` expectation has already returned
            Expired,
            Mut(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
            // Should be Box<dyn FnOnce> once that feature is stabilized
            // https://github.com/rust-lang/rust/issues/28796
            Once(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
        }

        impl  Rfunc {
            fn call_mut(&mut self, $( $args: $methty, )* ) -> $o {
                match self {
                    Rfunc::Default => {
                        unimplemented!()
                        //Self::return_default()
                    },
                    Rfunc::Expired => {
                        panic!("Called a method twice that was expected only once")
                    },
                    Rfunc::Mut(f) => {
                        f( $( $args, )* )
                    },
                    Rfunc::Once(_) => {
                        let fo = ::std::mem::replace(self, Rfunc::Expired);
                        if let Rfunc::Once(mut f) = fo {
                            f( $( $args, )* )
                        } else {
                            unreachable!()
                        }
                    },
                }
            }
        }

        impl std::default::Default for Rfunc {
            fn default() -> Self {
                Rfunc::Default
            }
        }

        enum Matcher {
            Func(Box<Fn($( &$matchty, )* ) -> bool>),
            Pred( $( Box<::mockall::Predicate<$matchty>>, )* )
        }

        impl Matcher {
            fn eval(&self, $( $args: &$matchty, )*) -> bool {
                match self {
                    Matcher::Func(f) => f($( $args, )*),
                    Matcher::Pred($( $predargs, )*) =>
                        [$( $predargs.eval($args), )*]
                        .into_iter()
                        .all(|x| *x),
                }
            }

            fn verify(&self, $( $args: &$matchty, )* ) {
                match self {
                    Matcher::Func(f) => assert!(f($( $args, )*),
                        "Expectation didn't match arguments"),
                    Matcher::Pred($( $predargs, )*) => {
                        $(if let Some(c) = $predargs.find_case(false, $args) {
                            panic!("Expectation didn't match arguments:\n{}",
                                   c.tree());
                        })*
                    }
                }
            }
        }

        impl Default for Matcher {
            #[allow(unused_variables)]
            fn default() -> Self {
                Matcher::Func(Box::new(|$( $args, )*| true))
            }
        }

        #[derive(Default)]
        struct Common {
            matcher: ::std::sync::Mutex<Matcher>,
            seq_handle: Option<::mockall::SeqHandle>,
            times: ::mockall::Times
        }

        impl Common {
            fn call(&self, $( $args: &$matchty, )* ) {
                self.matcher.lock().unwrap().verify($( $args, )*);
                self.times.call();
                self.verify_sequence();
                if self.times.is_satisfied() {
                    self.satisfy_sequence()
                }
            }

            fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                assert!(self.times.is_exact(),
                    "Only Expectations with an exact call count have sequences");
                self.seq_handle = Some(seq.next());
                self
            }

            /// Forbid this expectation from ever being called
            fn never(&mut self) {
                self.times.never();
            }

            fn satisfy_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.satisfy()
                }
            }

            /// Require this expectation to be called exactly `n` times.
            fn times(&mut self, n: usize) {
                self.times.n(n);
            }

            /// Allow this expectation to be called any number of times
            fn times_any(&mut self) {
                self.times.any();
            }

            /// Allow this expectation to be called any number of times within a given
            /// range
            fn times_range(&mut self, range: ::std::ops::Range<usize>) {
                self.times.range(range);
            }

            fn verify_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.verify()
                }
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*)
            {
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Pred($( Box::new($args), )*);
                ::std::mem::replace(guard.deref_mut(), m);
            }

            fn withf<F>(&mut self, f: F)
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Func(Box::new(f));
                ::std::mem::replace(guard.deref_mut(), m);
            }
        }

        #[derive(Default)]
        pub struct Expectation {
            common: Common,
            rfunc: ::std::sync::Mutex<Rfunc>,
        }

        impl Expectation {
            pub fn call(&self, $( $args: $methty, )* ) -> $o {
                self.common.call($( $matchcall, )*);
                self.rfunc.lock().unwrap().call_mut($( $args, )*)
            }

            pub fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                self.common.in_sequence(seq);
                self
            }

            /// Forbid this expectation from ever being called
            pub fn never(&mut self) -> &mut Self {
                self.common.never();
                self
            }

            /// Supply an `FnOnce` closure that will provide the return value
            /// for this Expectation.  This is useful for return types that
            /// aren't `Clone`.  It will be an error to call this Expectation
            /// multiple times.
            pub fn return_once<F>(&mut self, f: F) -> &mut Self
                where F: FnOnce($( $methty, )*) -> $o + Send + 'static
            {
                let mut fopt = Some(f);
                let fmut = move |$( $args: $methty, )*| {
                    if let Some(f) = fopt.take() {
                        f($( $args, )*)
                    } else {
                        panic!("Called a method twice that was expected only once")
                    }
                };
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(),
                                        Rfunc::Once(Box::new(fmut)));
                }
                self
            }

            /// Single-threaded version of [`return_once`](#method.return_once).
            /// This is useful for return types that are neither `Send` nor
            /// `Clone`.
            ///
            /// It is a runtime error to call the mock method from a different
            /// thread than the one that originally called this method.  It is
            /// also a runtime error to call the method more than once.
            pub fn return_once_st<F>(&mut self, f: F) -> &mut Self
                where F: FnOnce($( $methty, )*) -> $o + 'static
            {
                let mut fragile = Some(::fragile::Fragile::new(f));
                let fmut = Box::new(move |$( $args: $methty, )*| {
                    match fragile.take() {
                        Some(frag) => (frag.into_inner())($( $args, )*),
                        None => panic!(
                            "Called a method twice that was expected only once")
                    }
                });
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(), Rfunc::Once(fmut));
                }
                self
            }

            pub fn returning<F>(&mut self, f: F) -> &mut Self
                where F: FnMut($( $methty, )*) -> $o + Send + 'static
            {
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(), Rfunc::Mut(Box::new(f)));
                }
                self
            }

            /// Single-threaded version of [`returning`](#method.returning).
            /// Can be used when the argument or return type isn't `Send`.
            ///
            /// It is a runtime error to call the mock method from a different
            /// thread than the one that originally called this method.
            pub fn returning_st<F>(&mut self, f: F) -> &mut Self
                where F: FnMut($( $methty, )*) -> $o + 'static
            {
                let mut fragile = ::fragile::Fragile::new(f);
                let fmut = move |$( $args: $methty, )*| {
                    (fragile.get_mut())($( $args, )*)
                };
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(),
                                        Rfunc::Mut(Box::new(fmut)));
                }
                self
            }

            /// Require this expectation to be called exactly `n` times.
            pub fn times(&mut self, n: usize) -> &mut Self {
                self.common.times(n);
                self
            }

            /// Allow this expectation to be called any number of times
            ///
            /// This behavior is the default, but the method is provided in case the
            /// default behavior changes.
            pub fn times_any(&mut self) -> &mut Self {
                self.common.times_any();
                self
            }

            /// Allow this expectation to be called any number of times within a
            /// given range
            pub fn times_range(&mut self, range: ::std::ops::Range<usize>)
                -> &mut Self
            {
                self.common.times_range(range);
                self
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            pub fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*) -> &mut Self
            {
                self.common.with($( $args, )*);
                self
            }

            pub fn withf<F>(&mut self, f: F) -> &mut Self
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                self.common.withf(f);
                self
            }
        }

        }
        use $module::Expectation as $expectation;
    }
}
<|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 crate::types::{type_to_cpp, TypeName};
use itertools::Itertools;
use std::collections::HashSet;
use syn::{parse_quote, Ident, Type};

#[derive(Clone)]
enum ArgumentConversionType {
    None,
    FromUniquePtrToValue,
    FromValueToUniquePtr,
}

#[derive(Clone)]
pub(crate) struct ArgumentConversion {
    unwrapped_type: Type,
    conversion: ArgumentConversionType,
}

impl ArgumentConversion {
    pub(crate) fn new_unconverted(ty: Type) -> Self {
        ArgumentConversion {
            unwrapped_type: ty,
            conversion: ArgumentConversionType::None,
        }
    }

    pub(crate) fn new_to_unique_ptr(ty: Type) -> Self {
        ArgumentConversion {
            unwrapped_type: ty,
            conversion: ArgumentConversionType::FromValueToUniquePtr,
        }
    }

    pub(crate) fn new_from_unique_ptr(ty: Type) -> Self {
        ArgumentConversion {
            unwrapped_type: ty,
            conversion: ArgumentConversionType::FromUniquePtrToValue,
        }
    }

    pub(crate) fn work_needed(&self) -> bool {
        !matches!(self.conversion, ArgumentConversionType::None)
    }

    fn unconverted_type(&self) -> String {
        match self.conversion {
            ArgumentConversionType::FromUniquePtrToValue => self.wrapped_type(),
            _ => self.unwrapped_type_as_string(),
        }
    }

    fn converted_type(&self) -> String {
        match self.conversion {
            ArgumentConversionType::FromValueToUniquePtr => self.wrapped_type(),
            _ => self.unwrapped_type_as_string(),
        }
    }

    pub(crate) fn unconverted_rust_type(&self) -> Type {
        match self.conversion {
            ArgumentConversionType::FromValueToUniquePtr => {
                let orig_ty = &self.unwrapped_type;
                parse_quote! {
                    UniquePtr < #orig_ty >
                }
            }
            _ => self.unwrapped_type.clone(),
        }
    }

    pub(crate) fn converted_rust_type(&self) -> Type {
        match self.conversion {
            ArgumentConversionType::FromUniquePtrToValue => {
                let orig_ty = &self.unwrapped_type;
                parse_quote! {
                    UniquePtr < #orig_ty >
                }
            }
            _ => self.unwrapped_type.clone(),
        }
    }

    fn unwrapped_type_as_string(&self) -> String {
        type_to_cpp(&self.unwrapped_type, TypeName::from_cxx_type_path)
    }

    fn wrapped_type(&self) -> String {
        format!("std::unique_ptr<{}>", self.unwrapped_type_as_string())
    }

    fn conversion(&self, var_name: &str) -> String {
        match self.conversion {
            ArgumentConversionType::None => var_name.to_string(),
            ArgumentConversionType::FromUniquePtrToValue => format!("std::move(*{})", var_name),
            ArgumentConversionType::FromValueToUniquePtr => format!(
                "std::make_unique<{}>({})",
                self.unconverted_type(),
                var_name
            ),
        }
    }
}

pub(crate) struct ByValueWrapper {
    pub(crate) id: Ident,
    pub(crate) return_conversion: Option<ArgumentConversion>,
    pub(crate) argument_conversion: Vec<ArgumentConversion>,
    pub(crate) is_a_method: bool,
}

/// Instructions for new C++ which we need to generate.
pub(crate) enum AdditionalNeed {
    MakeStringConstructor,
    MakeUnique(TypeName, Vec<TypeName>),
    ByValueWrapper(Box<ByValueWrapper>),
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Hash)]
struct Header {
    name: &'static str,
    system: bool,
}

impl Header {
    fn system(name: &'static str) -> Self {
        Header { name, system: true }
    }

    fn user(name: &'static str) -> Self {
        Header {
            name,
            system: false,
        }
    }

    fn include_stmt(&self) -> String {
        if self.system {
            format!("#include <{}>", self.name)
        } else {
            format!("#include \"{}\"", self.name)
        }
    }
}

struct AdditionalFunction {
    declaration: String,
    definition: String,
    headers: Vec<Header>,
}

/// Details of additional generated C++.
pub(crate) struct AdditionalCpp {
    pub(crate) declarations: String,
    pub(crate) definitions: String,
}

/// Generates additional C++ glue functions needed by autocxx.
/// In some ways it would be preferable to be able to pass snippets
/// of C++ through to `cxx` for inclusion in the C++ file which it
/// generates, and perhaps we'll explore that in future. But for now,
/// autocxx generates its own _additional_ C++ files which therefore
/// need to be built and included in linking procedures.
pub(crate) struct AdditionalCppGenerator {
    additional_functions: Vec<AdditionalFunction>,
    inclusions: String,
}

impl AdditionalCppGenerator {
    pub(crate) fn new(inclusions: String) -> Self {
        AdditionalCppGenerator {
            additional_functions: Vec::new(),
            inclusions,
        }
    }

    pub(crate) fn add_needs(&mut self, additions: Vec<AdditionalNeed>) {
        for need in additions {
            match need {
                AdditionalNeed::MakeStringConstructor => self.generate_string_constructor(),
                AdditionalNeed::MakeUnique(ty, args) => self.generate_make_unique(&ty, &args),
                AdditionalNeed::ByValueWrapper(by_value_wrapper) => {
                    self.generate_by_value_wrapper(*by_value_wrapper)
                }
            }
        }
    }

    pub(crate) fn generate(&self) -> Option<AdditionalCpp> {
        if self.additional_functions.is_empty() {
            None
        } else {
            let headers: HashSet<Header> = self
                .additional_functions
                .iter()
                .map(|x| x.headers.iter().cloned())
                .flatten()
                .collect();
            let headers = headers.iter().map(|x| x.include_stmt()).join("\n");
            let declarations = self.concat_additional_items(|x| &x.declaration);
            let declarations = format!("{}\n{}\n{}", headers, self.inclusions, declarations);
            let definitions = self.concat_additional_items(|x| &x.definition);
            let definitions = format!("#include \"autocxxgen.h\"\n{}", definitions);
            Some(AdditionalCpp {
                declarations,
                definitions,
            })
        }
    }

    fn concat_additional_items<F>(&self, field_access: F) -> String
    where
        F: FnMut(&AdditionalFunction) -> &str,
    {
        let mut s = self
            .additional_functions
            .iter()
            .map(field_access)
            .collect::<Vec<&str>>()
            .join("\n\n");
        s.push('\n');
        s
    }

    fn generate_string_constructor(&mut self) {
        let declaration = "std::unique_ptr<std::string> make_string(::rust::Str str)";
        let definition = format!(
            "{} {{ return std::make_unique<std::string>(std::string(str)); }}",
            declaration
        );
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            declaration,
            definition,
            headers: vec![
                Header::system("memory"),
                Header::system("string"),
                Header::user("cxx.h"),
            ],
        })
    }

    fn generate_make_unique(&mut self, ty: &TypeName, constructor_arg_types: &[TypeName]) {
        let name = format!("{}_make_unique", ty.to_cpp_name());
        let constructor_args = constructor_arg_types
            .iter()
            .enumerate()
            .map(|(counter, ty)| format!("{} arg{}", ty.to_cpp_name(), counter))
            .join(", ");
        let declaration = format!("std::unique_ptr<{}> {}({})", ty, name, constructor_args);
        let arg_list = constructor_arg_types
            .iter()
            .enumerate()
            .map(|(counter, _)| format!("arg{}", counter))
            .join(", ");
        let definition = format!(
            "{} {{ return std::make_unique<{}>({}); }}",
            declaration, ty, arg_list
        );
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            declaration,
            definition,
            headers: vec![Header::system("memory")],
        })
    }

    fn generate_by_value_wrapper(&mut self, details: ByValueWrapper) {
        let ident = details.id;
        let is_a_method = details.is_a_method;
        let name = format!("{}_up_wrapper", ident.to_string());
        let get_arg_name = |counter: usize| -> String {
            if is_a_method && counter == 0 {
                // For method calls that we generate, the first
                // argument name needs to be such that we recognize
                // it as a method in the second invocation of
                // bridge_converter after it's flowed again through
                // bindgen.
                "autocxx_gen_this".to_string()
            } else {
                format!("arg{}", counter)
            }
        };
        let args = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, ty)| format!("{} {}", ty.unconverted_type(), get_arg_name(counter)))
            .join(", ");
        let ret_type = details
            .return_conversion
            .as_ref()
            .map_or("void".to_string(), |x| x.converted_type());
        let declaration = format!("{} {}({})", ret_type, name, args);
        let mut arg_list = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, conv)| conv.conversion(&get_arg_name(counter)));
        let receiver = if is_a_method { arg_list.next() } else { None };
        let arg_list = arg_list.join(", ");
        let mut underlying_function_call = format!("{}({})", ident.to_string(), arg_list);
        if let Some(receiver) = receiver {
            underlying_function_call = format!("{}.{}", receiver, underlying_function_call);
        }
        if let Some(ret) = details.return_conversion {
            underlying_function_call =
                format!("return {}", ret.conversion(&underlying_function_call));
        };
        let definition = format!("{} {{ {}; }}", declaration, underlying_function_call,);
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            declaration,
            definition,
            headers: vec![Header::system("memory")],
        })
    }
}
<commit_msg>Fix bug with namespaces in UP wrappers<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 crate::types::{type_to_cpp, TypeName};
use itertools::Itertools;
use std::collections::HashSet;
use syn::{parse_quote, Ident, Type};

#[derive(Clone)]
enum ArgumentConversionType {
    None,
    FromUniquePtrToValue,
    FromValueToUniquePtr,
}

#[derive(Clone)]
pub(crate) struct ArgumentConversion {
    unwrapped_type: Type,
    conversion: ArgumentConversionType,
}

impl ArgumentConversion {
    pub(crate) fn new_unconverted(ty: Type) -> Self {
        ArgumentConversion {
            unwrapped_type: ty,
            conversion: ArgumentConversionType::None,
        }
    }

    pub(crate) fn new_to_unique_ptr(ty: Type) -> Self {
        ArgumentConversion {
            unwrapped_type: ty,
            conversion: ArgumentConversionType::FromValueToUniquePtr,
        }
    }

    pub(crate) fn new_from_unique_ptr(ty: Type) -> Self {
        ArgumentConversion {
            unwrapped_type: ty,
            conversion: ArgumentConversionType::FromUniquePtrToValue,
        }
    }

    pub(crate) fn work_needed(&self) -> bool {
        !matches!(self.conversion, ArgumentConversionType::None)
    }

    fn unconverted_type(&self) -> String {
        match self.conversion {
            ArgumentConversionType::FromUniquePtrToValue => self.wrapped_type(),
            _ => self.unwrapped_type_as_string(),
        }
    }

    fn converted_type(&self) -> String {
        match self.conversion {
            ArgumentConversionType::FromValueToUniquePtr => self.wrapped_type(),
            _ => self.unwrapped_type_as_string(),
        }
    }

    pub(crate) fn unconverted_rust_type(&self) -> Type {
        match self.conversion {
            ArgumentConversionType::FromValueToUniquePtr => self.make_unique_ptr_type(),
            _ => self.unwrapped_type.clone(),
        }
    }

    pub(crate) fn converted_rust_type(&self) -> Type {
        match self.conversion {
            ArgumentConversionType::FromUniquePtrToValue => self.make_unique_ptr_type(),
            _ => self.unwrapped_type.clone(),
        }
    }

    fn unwrapped_type_as_string(&self) -> String {
        type_to_cpp(&self.unwrapped_type, TypeName::from_cxx_type_path)
    }

    fn wrapped_type(&self) -> String {
        format!("std::unique_ptr<{}>", self.unwrapped_type_as_string())
    }

    fn conversion(&self, var_name: &str) -> String {
        match self.conversion {
            ArgumentConversionType::None => var_name.to_string(),
            ArgumentConversionType::FromUniquePtrToValue => format!("std::move(*{})", var_name),
            ArgumentConversionType::FromValueToUniquePtr => format!(
                "std::make_unique<{}>({})",
                self.unconverted_type(),
                var_name
            ),
        }
    }

    fn make_unique_ptr_type(&self) -> Type {
        let innerty = match &self.unwrapped_type {
            Type::Path(typ) => {
                // Until cxx supports a hierarchic set of inner mods
                // for namespace purposes, we just take the final segment.
                let final_seg = typ.path.segments.last().unwrap();
                parse_quote! {
                   #final_seg
                }
            }
            _ => self.unwrapped_type.clone(),
        };
        parse_quote! {
            UniquePtr < #innerty >
        }
    }
}

pub(crate) struct ByValueWrapper {
    pub(crate) id: Ident,
    pub(crate) return_conversion: Option<ArgumentConversion>,
    pub(crate) argument_conversion: Vec<ArgumentConversion>,
    pub(crate) is_a_method: bool,
}

/// Instructions for new C++ which we need to generate.
pub(crate) enum AdditionalNeed {
    MakeStringConstructor,
    MakeUnique(TypeName, Vec<TypeName>),
    ByValueWrapper(Box<ByValueWrapper>),
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Hash)]
struct Header {
    name: &'static str,
    system: bool,
}

impl Header {
    fn system(name: &'static str) -> Self {
        Header { name, system: true }
    }

    fn user(name: &'static str) -> Self {
        Header {
            name,
            system: false,
        }
    }

    fn include_stmt(&self) -> String {
        if self.system {
            format!("#include <{}>", self.name)
        } else {
            format!("#include \"{}\"", self.name)
        }
    }
}

struct AdditionalFunction {
    declaration: String,
    definition: String,
    headers: Vec<Header>,
}

/// Details of additional generated C++.
pub(crate) struct AdditionalCpp {
    pub(crate) declarations: String,
    pub(crate) definitions: String,
}

/// Generates additional C++ glue functions needed by autocxx.
/// In some ways it would be preferable to be able to pass snippets
/// of C++ through to `cxx` for inclusion in the C++ file which it
/// generates, and perhaps we'll explore that in future. But for now,
/// autocxx generates its own _additional_ C++ files which therefore
/// need to be built and included in linking procedures.
pub(crate) struct AdditionalCppGenerator {
    additional_functions: Vec<AdditionalFunction>,
    inclusions: String,
}

impl AdditionalCppGenerator {
    pub(crate) fn new(inclusions: String) -> Self {
        AdditionalCppGenerator {
            additional_functions: Vec::new(),
            inclusions,
        }
    }

    pub(crate) fn add_needs(&mut self, additions: Vec<AdditionalNeed>) {
        for need in additions {
            match need {
                AdditionalNeed::MakeStringConstructor => self.generate_string_constructor(),
                AdditionalNeed::MakeUnique(ty, args) => self.generate_make_unique(&ty, &args),
                AdditionalNeed::ByValueWrapper(by_value_wrapper) => {
                    self.generate_by_value_wrapper(*by_value_wrapper)
                }
            }
        }
    }

    pub(crate) fn generate(&self) -> Option<AdditionalCpp> {
        if self.additional_functions.is_empty() {
            None
        } else {
            let headers: HashSet<Header> = self
                .additional_functions
                .iter()
                .map(|x| x.headers.iter().cloned())
                .flatten()
                .collect();
            let headers = headers.iter().map(|x| x.include_stmt()).join("\n");
            let declarations = self.concat_additional_items(|x| &x.declaration);
            let declarations = format!("{}\n{}\n{}", headers, self.inclusions, declarations);
            let definitions = self.concat_additional_items(|x| &x.definition);
            let definitions = format!("#include \"autocxxgen.h\"\n{}", definitions);
            Some(AdditionalCpp {
                declarations,
                definitions,
            })
        }
    }

    fn concat_additional_items<F>(&self, field_access: F) -> String
    where
        F: FnMut(&AdditionalFunction) -> &str,
    {
        let mut s = self
            .additional_functions
            .iter()
            .map(field_access)
            .collect::<Vec<&str>>()
            .join("\n\n");
        s.push('\n');
        s
    }

    fn generate_string_constructor(&mut self) {
        let declaration = "std::unique_ptr<std::string> make_string(::rust::Str str)";
        let definition = format!(
            "{} {{ return std::make_unique<std::string>(std::string(str)); }}",
            declaration
        );
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            declaration,
            definition,
            headers: vec![
                Header::system("memory"),
                Header::system("string"),
                Header::user("cxx.h"),
            ],
        })
    }

    fn generate_make_unique(&mut self, ty: &TypeName, constructor_arg_types: &[TypeName]) {
        let name = format!("{}_make_unique", ty.to_cpp_name());
        let constructor_args = constructor_arg_types
            .iter()
            .enumerate()
            .map(|(counter, ty)| format!("{} arg{}", ty.to_cpp_name(), counter))
            .join(", ");
        let declaration = format!("std::unique_ptr<{}> {}({})", ty, name, constructor_args);
        let arg_list = constructor_arg_types
            .iter()
            .enumerate()
            .map(|(counter, _)| format!("arg{}", counter))
            .join(", ");
        let definition = format!(
            "{} {{ return std::make_unique<{}>({}); }}",
            declaration, ty, arg_list
        );
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            declaration,
            definition,
            headers: vec![Header::system("memory")],
        })
    }

    fn generate_by_value_wrapper(&mut self, details: ByValueWrapper) {
        let ident = details.id;
        let is_a_method = details.is_a_method;
        let name = format!("{}_up_wrapper", ident.to_string());
        let get_arg_name = |counter: usize| -> String {
            if is_a_method && counter == 0 {
                // For method calls that we generate, the first
                // argument name needs to be such that we recognize
                // it as a method in the second invocation of
                // bridge_converter after it's flowed again through
                // bindgen.
                "autocxx_gen_this".to_string()
            } else {
                format!("arg{}", counter)
            }
        };
        let args = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, ty)| format!("{} {}", ty.unconverted_type(), get_arg_name(counter)))
            .join(", ");
        let ret_type = details
            .return_conversion
            .as_ref()
            .map_or("void".to_string(), |x| x.converted_type());
        let declaration = format!("{} {}({})", ret_type, name, args);
        let mut arg_list = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, conv)| conv.conversion(&get_arg_name(counter)));
        let receiver = if is_a_method { arg_list.next() } else { None };
        let arg_list = arg_list.join(", ");
        let mut underlying_function_call = format!("{}({})", ident.to_string(), arg_list);
        if let Some(receiver) = receiver {
            underlying_function_call = format!("{}.{}", receiver, underlying_function_call);
        }
        if let Some(ret) = details.return_conversion {
            underlying_function_call =
                format!("return {}", ret.conversion(&underlying_function_call));
        };
        let definition = format!("{} {{ {}; }}", declaration, underlying_function_call,);
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            declaration,
            definition,
            headers: vec![Header::system("memory")],
        })
    }
}
<|endoftext|>