<commit_before>use std::collections::HashSet;

use crate::{Assist, AssistId, AssistCtx};

use hir::Resolver;
use hir::db::HirDatabase;
use ra_syntax::{SmolStr, SyntaxKind, TextRange, TextUnit, TreeArc};
use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
use ra_db::FilePosition;
use ra_fmt::{leading_indent, reindent};

use itertools::Itertools;

/// Given an `ast::ImplBlock`, resolves the target trait (the one being
/// implemented) to a `ast::TraitDef`.
pub(crate) fn resolve_target_trait_def(
    db: &impl HirDatabase,
    resolver: &Resolver,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path = impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?;
    let hir_path = ast_path.path().and_then(hir::Path::from_ast)?;

    match resolver.resolve_path(db, &hir_path).take_types() {
        Some(hir::Resolution::Def(hir::ModuleDef::Trait(def))) => Some(def.source(db).1),
        _ => None,
    }
}

pub(crate) fn build_func_body(def: &ast::FnDef) -> String {
    let mut buf = String::new();

    for child in def.syntax().children() {
        if child.kind() == SyntaxKind::SEMI {
            buf.push_str(" { unimplemented!() }")
        } else {
            child.text().push_to(&mut buf);
        }
    }

    buf.trim_end().to_string()
}

pub(crate) fn add_missing_impl_members(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let node = ctx.covering_node();
    let impl_node = node.ancestors().find_map(ast::ImplBlock::cast)?;
    let impl_item_list = impl_node.item_list()?;
    // Don't offer the assist when cursor is at the end, outside the block itself.
    if node.range().end() == impl_node.syntax().range().end() {
        return None;
    }

    let trait_def = {
        let db = ctx.db;
        // TODO: Can we get the position of cursor itself rather than supplied range?
        let range = ctx.frange;
        let position = FilePosition { file_id: range.file_id, offset: range.range.start() };
        let resolver = hir::source_binder::resolver_for_position(db, position);

        resolve_target_trait_def(db, &resolver, impl_node)?
    };

    let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
    let def_name = |def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };

    let trait_items = trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
    let impl_items = impl_item_list.impl_items();

    let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
    let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();

    let trait_fn_names = trait_fns.iter().cloned().filter_map(def_name).collect::<HashSet<_>>();
    let impl_fn_names = impl_fns.iter().cloned().filter_map(def_name).collect::<HashSet<_>>();

    let missing_fn_names = trait_fn_names.difference(&impl_fn_names).collect::<HashSet<_>>();
    let missing_fns: Vec<_> = trait_fns
        .iter()
        .cloned()
        .filter(|t| def_name(t).map(|n| missing_fn_names.contains(&n)).unwrap_or(false))
        .collect();
    if missing_fns.is_empty() {
        return None;
    }

    ctx.add_action(AssistId("add_impl_missing_members"), "add missing impl members", |edit| {
        let indent = {
            // FIXME: Find a way to get the indent already used in the file.
            // Now, we copy the indent of first item or indent with 4 spaces relative to impl block
            const DEFAULT_INDENT: &str = "    ";
            let first_item = impl_item_list.impl_items().next();
            let first_item_indent = first_item.and_then(|i| leading_indent(i.syntax()));
            let impl_block_indent = || leading_indent(impl_node.syntax()).unwrap_or_default();

            first_item_indent
                .map(ToOwned::to_owned)
                .unwrap_or_else(|| impl_block_indent().to_owned() + DEFAULT_INDENT)
        };

        let mut func_bodies = missing_fns.into_iter().map(build_func_body);
        let func_bodies = func_bodies.join("\n");
        let func_bodies = String::from("\n") + &func_bodies;
        let func_bodies = reindent(&func_bodies, &indent) + "\n";

        let changed_range = {
            let last_whitespace = impl_item_list.syntax().children();
            let last_whitespace = last_whitespace.filter_map(ast::Whitespace::cast).last();
            let last_whitespace = last_whitespace.map(|w| w.syntax());

            let cursor_range = TextRange::from_to(node.range().end(), node.range().end());

            last_whitespace.map(|x| x.range()).unwrap_or(cursor_range)
        };

        let replaced_text_range = TextUnit::of_str(&func_bodies);

        edit.replace(changed_range, func_bodies);
        edit.set_cursor(changed_range.start() + replaced_text_range - TextUnit::of_str("\n"));
    });

    ctx.build()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable};

    #[test]
    fn test_add_missing_impl_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }
    fn baz(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {<|>}",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_cursor_after_empty_impl_block() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {}<|>",
        )
    }
}
<commit_msg>Simplify calculation of missing functions<commit_after>use crate::{Assist, AssistId, AssistCtx};

use hir::Resolver;
use hir::db::HirDatabase;
use ra_syntax::{SmolStr, SyntaxKind, TextRange, TextUnit, TreeArc};
use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
use ra_db::FilePosition;
use ra_fmt::{leading_indent, reindent};

use itertools::Itertools;

/// Given an `ast::ImplBlock`, resolves the target trait (the one being
/// implemented) to a `ast::TraitDef`.
pub(crate) fn resolve_target_trait_def(
    db: &impl HirDatabase,
    resolver: &Resolver,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path = impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?;
    let hir_path = ast_path.path().and_then(hir::Path::from_ast)?;

    match resolver.resolve_path(db, &hir_path).take_types() {
        Some(hir::Resolution::Def(hir::ModuleDef::Trait(def))) => Some(def.source(db).1),
        _ => None,
    }
}

pub(crate) fn build_func_body(def: &ast::FnDef) -> String {
    let mut buf = String::new();

    for child in def.syntax().children() {
        if child.kind() == SyntaxKind::SEMI {
            buf.push_str(" { unimplemented!() }")
        } else {
            child.text().push_to(&mut buf);
        }
    }

    buf.trim_end().to_string()
}

pub(crate) fn add_missing_impl_members(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let node = ctx.covering_node();
    let impl_node = node.ancestors().find_map(ast::ImplBlock::cast)?;
    let impl_item_list = impl_node.item_list()?;
    // Don't offer the assist when cursor is at the end, outside the block itself.
    if node.range().end() == impl_node.syntax().range().end() {
        return None;
    }

    let trait_def = {
        let db = ctx.db;
        // TODO: Can we get the position of cursor itself rather than supplied range?
        let range = ctx.frange;
        let position = FilePosition { file_id: range.file_id, offset: range.range.start() };
        let resolver = hir::source_binder::resolver_for_position(db, position);

        resolve_target_trait_def(db, &resolver, impl_node)?
    };

    let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
    let def_name = |def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };

    let trait_items = trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
    let impl_items = impl_item_list.impl_items();

    let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
    let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();

    let missing_fns: Vec<_> = trait_fns
        .into_iter()
        .filter(|t| impl_fns.iter().all(|i| def_name(i) != def_name(t)))
        .collect();
    if missing_fns.is_empty() {
        return None;
    }

    ctx.add_action(AssistId("add_impl_missing_members"), "add missing impl members", |edit| {
        let indent = {
            // FIXME: Find a way to get the indent already used in the file.
            // Now, we copy the indent of first item or indent with 4 spaces relative to impl block
            const DEFAULT_INDENT: &str = "    ";
            let first_item = impl_item_list.impl_items().next();
            let first_item_indent = first_item.and_then(|i| leading_indent(i.syntax()));
            let impl_block_indent = || leading_indent(impl_node.syntax()).unwrap_or_default();

            first_item_indent
                .map(ToOwned::to_owned)
                .unwrap_or_else(|| impl_block_indent().to_owned() + DEFAULT_INDENT)
        };

        let mut func_bodies = missing_fns.into_iter().map(build_func_body);
        let func_bodies = func_bodies.join("\n");
        let func_bodies = String::from("\n") + &func_bodies;
        let func_bodies = reindent(&func_bodies, &indent) + "\n";

        let changed_range = {
            let last_whitespace = impl_item_list.syntax().children();
            let last_whitespace = last_whitespace.filter_map(ast::Whitespace::cast).last();
            let last_whitespace = last_whitespace.map(|w| w.syntax());

            let cursor_range = TextRange::from_to(node.range().end(), node.range().end());

            last_whitespace.map(|x| x.range()).unwrap_or(cursor_range)
        };

        let replaced_text_range = TextUnit::of_str(&func_bodies);

        edit.replace(changed_range, func_bodies);
        edit.set_cursor(changed_range.start() + replaced_text_range - TextUnit::of_str("\n"));
    });

    ctx.build()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable};

    #[test]
    fn test_add_missing_impl_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }
    fn baz(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {<|>}",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_cursor_after_empty_impl_block() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {}<|>",
        )
    }
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate openssl;

use errors::common::CommonError;

use services::anoncreds::constants::LARGE_MVECT;
use utils::crypto::bn::BigNumber;
use utils::crypto::pair::GroupOrderElement;
use std::hash::Hash;
use std::cmp::max;
use std::collections::HashMap;

#[cfg(not(test))]
pub fn random_qr(n: &BigNumber) -> Result<BigNumber, CommonError> {
    let random = n
        .rand_range()?
        .sqr(None)?
        .modulus(&n, None)?;
    Ok(random)
}

#[cfg(test)]
pub fn random_qr(n: &BigNumber) -> Result<BigNumber, CommonError> {
    Ok(BigNumber::from_dec("64684820421150545443421261645532741305438158267230326415141505826951816460650437611148133267480407958360035501128469885271549378871140475869904030424615175830170939416512594291641188403335834762737251794282186335118831803135149622404791467775422384378569231649224208728902565541796896860352464500717052768431523703881746487372385032277847026560711719065512366600220045978358915680277126661923892187090579302197390903902744925313826817940566429968987709582805451008234648959429651259809188953915675063700676546393568304468609062443048457324721450190021552656280473128156273976008799243162970386898307404395608179975243")?)
}

pub fn bitwise_or_big_int(a: &BigNumber, b: &BigNumber) -> Result<BigNumber, CommonError> {
    let significant_bits = max(a.num_bits()?, b.num_bits()?);
    let mut result = BigNumber::new()?;
    for i in 0..significant_bits {
        if a.is_bit_set(i)? || b.is_bit_set(i)? {
            result.set_bit(i)?;
        }
    }
    Ok(result)
}

pub fn transform_u32_to_array_of_u8(x: u32) -> Vec<u8> {
    let mut result: Vec<u8> = vec![0; 60];
    for i in (0..4).rev() {
        let shift = i * 8;
        let b = (x >> shift) as u8;
        result.push(b);
    }
    result
}

pub fn get_hash_as_int(nums: &mut Vec<Vec<u8>>) -> Result<BigNumber, CommonError> {
    nums.sort();

    let mut hashed_array: Vec<u8> = BigNumber::hash_array(&nums)?;
    hashed_array.reverse();

    BigNumber::from_bytes(&hashed_array[..])
}

pub fn get_mtilde(unrevealed_attrs: &Vec<String>)
                  -> Result<HashMap<String, BigNumber>, CommonError> {
    let mut mtilde: HashMap<String, BigNumber> = HashMap::new();

    for attr in unrevealed_attrs.iter() {
        mtilde.insert(attr.clone(), BigNumber::rand(LARGE_MVECT)?);
    }
    Ok(mtilde)
}

fn largest_square_less_than(delta: i32) -> i32 {
    (delta as f64).sqrt().floor() as i32
}

pub fn four_squares(delta: i32) -> Result<HashMap<String, BigNumber>, CommonError> {
    let u1 = largest_square_less_than(delta);
    let u2 = largest_square_less_than(delta - u1.pow(2));
    let u3 = largest_square_less_than(delta - u1.pow(2) - u2.pow(2));
    let u4 = largest_square_less_than(delta - u1.pow(2) - u2.pow(2) - u3.pow(2));

    if u1.pow(2) + u2.pow(2) + u3.pow(2) + u4.pow(2) == delta {
        let mut res: HashMap<String, BigNumber> = HashMap::new();
        res.insert("0".to_string(), BigNumber::from_dec(&u1.to_string()[..])?);
        res.insert("1".to_string(), BigNumber::from_dec(&u2.to_string()[..])?);
        res.insert("2".to_string(), BigNumber::from_dec(&u3.to_string()[..])?);
        res.insert("3".to_string(), BigNumber::from_dec(&u4.to_string()[..])?);

        Ok(res)
    } else {
        Err(CommonError::InvalidStructure(format!("Cannot get the four squares for delta {} ", delta)))
    }
}

pub trait BytesView {
    fn to_bytes(&self) -> Result<Vec<u8>, CommonError>;
}

pub trait AppendByteArray {
    fn append_vec<T: BytesView>(&mut self, other: &Vec<T>) -> Result<(), CommonError>;
}

impl AppendByteArray for Vec<Vec<u8>> {
    fn append_vec<T: BytesView>(&mut self, other: &Vec<T>) -> Result<(), CommonError> {
        for el in other.iter() {
            self.push(el.to_bytes()?);
        }
        Ok(())
    }
}

pub fn clone_bignum_map<K: Clone + Eq + Hash>(other: &HashMap<K, BigNumber>)
                                              -> Result<HashMap<K, BigNumber>, CommonError> {
    let mut res: HashMap<K, BigNumber> = HashMap::new();
    for (k, v) in other {
        res.insert(k.clone(), v.clone()?);
    }
    Ok(res)
}

pub fn group_element_to_bignum(el: &GroupOrderElement) -> Result<BigNumber, CommonError> {
    Ok(BigNumber::from_bytes(&el.to_bytes()?)?)
}

pub fn bignum_to_group_element(num: &BigNumber) -> Result<GroupOrderElement, CommonError> {
    Ok(GroupOrderElement::from_bytes(&num.to_bytes()?)?)
}

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

    #[test]
    fn bitwise_or_big_int_works() {
        let a = BigNumber::from_dec("778378032744961463933002553964902776831187587689736807008034459507677878432383414623740074");
        let b = BigNumber::from_dec("1018517988167243043134222844204689080525734196832968125318070224677190649881668353091698688");
        let result = BigNumber::from_dec("1796896020912204507067225398169591857356921784522704932326104684184868528314051767715438762");
        assert_eq!(result.unwrap(), bitwise_or_big_int(&a.unwrap(), &b.unwrap()).unwrap());
    }

    #[test]
    fn get_hash_as_int_works() {
        let mut nums = vec![
            BigNumber::from_hex("ff9d2eedfee9cffd9ef6dbffedff3fcbef4caecb9bffe79bfa94d3fdf6abfbff").unwrap().to_bytes().unwrap(),
            BigNumber::from_hex("ff9d2eedfee9cffd9ef6dbffedff3fcbef4caecb9bffe79bfa9168615ccbc546").unwrap().to_bytes().unwrap()
        ];
        let res = get_hash_as_int(&mut nums);

        assert!(res.is_ok());
        assert_eq!("9E2A0653691B96A9B55B3D1133F9FEE2F2C37B848DBADF2F70DFFFE9E47C5A5D", res.unwrap().to_hex().unwrap());
    }

    #[test]
    fn four_squares_works() {
        let res = four_squares(10 as i32);

        assert!(res.is_ok());
        let res_data = res.unwrap();

        assert_eq!("3".to_string(), res_data.get("0").unwrap().to_dec().unwrap());
        assert_eq!("1".to_string(), res_data.get("1").unwrap().to_dec().unwrap());
        assert_eq!("0".to_string(), res_data.get("2").unwrap().to_dec().unwrap());
        assert_eq!("0".to_string(), res_data.get("3").unwrap().to_dec().unwrap());
    }

    #[test]
    fn transform_u32_to_array_of_u8_works() {
        let int = 1958376517;
        let answer = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 116, 186, 116, 69];
        assert_eq!(transform_u32_to_array_of_u8(int), answer)
    }
}<commit_msg>fixed test<commit_after>extern crate rand;
extern crate openssl;

use errors::common::CommonError;

use services::anoncreds::constants::LARGE_MVECT;
use utils::crypto::bn::BigNumber;
use utils::crypto::pair::GroupOrderElement;
use std::hash::Hash;
use std::cmp::max;
use std::collections::HashMap;

#[cfg(not(test))]
pub fn random_qr(n: &BigNumber) -> Result<BigNumber, CommonError> {
    let random = n
        .rand_range()?
        .sqr(None)?
        .modulus(&n, None)?;
    Ok(random)
}

#[cfg(test)]
pub fn random_qr(n: &BigNumber) -> Result<BigNumber, CommonError> {
    Ok(BigNumber::from_dec("64684820421150545443421261645532741305438158267230326415141505826951816460650437611148133267480407958360035501128469885271549378871140475869904030424615175830170939416512594291641188403335834762737251794282186335118831803135149622404791467775422384378569231649224208728902565541796896860352464500717052768431523703881746487372385032277847026560711719065512366600220045978358915680277126661923892187090579302197390903902744925313826817940566429968987709582805451008234648959429651259809188953915675063700676546393568304468609062443048457324721450190021552656280473128156273976008799243162970386898307404395608179975243")?)
}

pub fn bitwise_or_big_int(a: &BigNumber, b: &BigNumber) -> Result<BigNumber, CommonError> {
    let significant_bits = max(a.num_bits()?, b.num_bits()?);
    let mut result = BigNumber::new()?;
    for i in 0..significant_bits {
        if a.is_bit_set(i)? || b.is_bit_set(i)? {
            result.set_bit(i)?;
        }
    }
    Ok(result)
}

pub fn transform_u32_to_array_of_u8(x: u32) -> Vec<u8> {
    let mut result: Vec<u8> = vec![0; 60];
    for i in (0..4).rev() {
        let shift = i * 8;
        let b = (x >> shift) as u8;
        result.push(b);
    }
    result
}

pub fn get_hash_as_int(nums: &mut Vec<Vec<u8>>) -> Result<BigNumber, CommonError> {
    nums.sort();

    let mut hashed_array: Vec<u8> = BigNumber::hash_array(&nums)?;
    hashed_array.reverse();

    BigNumber::from_bytes(&hashed_array[..])
}

pub fn get_mtilde(unrevealed_attrs: &Vec<String>)
                  -> Result<HashMap<String, BigNumber>, CommonError> {
    let mut mtilde: HashMap<String, BigNumber> = HashMap::new();

    for attr in unrevealed_attrs.iter() {
        mtilde.insert(attr.clone(), BigNumber::rand(LARGE_MVECT)?);
    }
    Ok(mtilde)
}

fn largest_square_less_than(delta: i32) -> i32 {
    (delta as f64).sqrt().floor() as i32
}

pub fn four_squares(delta: i32) -> Result<HashMap<String, BigNumber>, CommonError> {
    let u1 = largest_square_less_than(delta);
    let u2 = largest_square_less_than(delta - u1.pow(2));
    let u3 = largest_square_less_than(delta - u1.pow(2) - u2.pow(2));
    let u4 = largest_square_less_than(delta - u1.pow(2) - u2.pow(2) - u3.pow(2));

    if u1.pow(2) + u2.pow(2) + u3.pow(2) + u4.pow(2) == delta {
        let mut res: HashMap<String, BigNumber> = HashMap::new();
        res.insert("0".to_string(), BigNumber::from_dec(&u1.to_string()[..])?);
        res.insert("1".to_string(), BigNumber::from_dec(&u2.to_string()[..])?);
        res.insert("2".to_string(), BigNumber::from_dec(&u3.to_string()[..])?);
        res.insert("3".to_string(), BigNumber::from_dec(&u4.to_string()[..])?);

        Ok(res)
    } else {
        Err(CommonError::InvalidStructure(format!("Cannot get the four squares for delta {} ", delta)))
    }
}

pub trait BytesView {
    fn to_bytes(&self) -> Result<Vec<u8>, CommonError>;
}

pub trait AppendByteArray {
    fn append_vec<T: BytesView>(&mut self, other: &Vec<T>) -> Result<(), CommonError>;
}

impl AppendByteArray for Vec<Vec<u8>> {
    fn append_vec<T: BytesView>(&mut self, other: &Vec<T>) -> Result<(), CommonError> {
        for el in other.iter() {
            self.push(el.to_bytes()?);
        }
        Ok(())
    }
}

pub fn clone_bignum_map<K: Clone + Eq + Hash>(other: &HashMap<K, BigNumber>)
                                              -> Result<HashMap<K, BigNumber>, CommonError> {
    let mut res: HashMap<K, BigNumber> = HashMap::new();
    for (k, v) in other {
        res.insert(k.clone(), v.clone()?);
    }
    Ok(res)
}

pub fn group_element_to_bignum(el: &GroupOrderElement) -> Result<BigNumber, CommonError> {
    Ok(BigNumber::from_bytes(&el.to_bytes()?)?)
}

pub fn bignum_to_group_element(num: &BigNumber) -> Result<GroupOrderElement, CommonError> {
    Ok(GroupOrderElement::from_bytes(&num.to_bytes()?)?)
}

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

    #[test]
    fn bitwise_or_big_int_works() {
        let a = BigNumber::from_dec("778378032744961463933002553964902776831187587689736807008034459507677878432383414623740074");
        let b = BigNumber::from_dec("1018517988167243043134222844204689080525734196832968125318070224677190649881668353091698688");
        let result = BigNumber::from_dec("1796896020912204507067225398169591857356921784522704932326104684184868528314051767715438762");
        assert_eq!(result.unwrap(), bitwise_or_big_int(&a.unwrap(), &b.unwrap()).unwrap());
    }

    #[test]
    fn get_hash_as_int_works() {
        let mut nums = vec![
            BigNumber::from_hex("ff9d2eedfee9cffd9ef6dbffedff3fcbef4caecb9bffe79bfa94d3fdf6abfbff").unwrap().to_bytes().unwrap(),
            BigNumber::from_hex("ff9d2eedfee9cffd9ef6dbffedff3fcbef4caecb9bffe79bfa9168615ccbc546").unwrap().to_bytes().unwrap()
        ];
        let res = get_hash_as_int(&mut nums);

        assert!(res.is_ok());
        assert_eq!("9E2A0653691B96A9B55B3D1133F9FEE2F2C37B848DBADF2F70DFFFE9E47C5A5D", res.unwrap().to_hex().unwrap());
    }

    #[test]
    fn four_squares_works() {
        let res = four_squares(10 as i32);

        assert!(res.is_ok());
        let res_data = res.unwrap();

        assert_eq!("3".to_string(), res_data.get("0").unwrap().to_dec().unwrap());
        assert_eq!("1".to_string(), res_data.get("1").unwrap().to_dec().unwrap());
        assert_eq!("0".to_string(), res_data.get("2").unwrap().to_dec().unwrap());
        assert_eq!("0".to_string(), res_data.get("3").unwrap().to_dec().unwrap());
    }

    #[test]
    fn transform_u32_to_array_of_u8_works() {
        let int = 1958376517;
        let answer = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 116, 186, 116, 69];
        assert_eq!(transform_u32_to_array_of_u8(int), answer)
    }
}<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{PerfLogEvent, PerfLogger};
use fnv::FnvHashSet;
use graphql_ir::{Program, ValidationResult};
use graphql_transforms::{
    apply_fragment_arguments, client_extensions, dedupe_type_discriminator, flatten,
    generate_id_field, generate_live_query_metadata, generate_preloadable_metadata,
    generate_subscription_name_metadata, generate_typename, handle_field_transform,
    inline_data_fragment, inline_fragments, mask, relay_early_flush, remove_base_fragments,
    skip_client_extensions, skip_redundant_nodes, skip_split_operation, skip_unreachable_node,
    skip_unused_variables, split_module_import, transform_connections, transform_defer_stream,
    transform_match, transform_refetchable_fragment, unwrap_custom_directive_selection,
    validate_module_conflicts, ConnectionInterface,
};
use interner::StringKey;
use std::sync::Arc;

pub struct Programs {
    pub source: Arc<Program>,
    pub reader: Arc<Program>,
    pub normalization: Arc<Program>,
    pub operation_text: Arc<Program>,
    pub typegen: Arc<Program>,
}

pub fn apply_transforms<TPerfLogger>(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    connection_interface: Arc<ConnectionInterface>,
    perf_logger: Arc<TPerfLogger>,
) -> ValidationResult<Programs>
where
    TPerfLogger: PerfLogger + 'static,
{
    // The execution pipeline is as follows, where items at the same indentation
    // can be computed independently and therefore in parallel:
    // |- common
    //  |- reader
    //  |- operation
    //     |- normalization
    //     |- operation_text
    // |- typegen
    //
    // NOTE: try_join(f1, f2) prefers the errors from f1 over f2, so process the normalization
    // program first since it is likely to include more errors (both per-fragment and
    // whole-operation errors) than the reader_program which just includes per-fragment errors.
    let (((normalization_program, text_program), reader_program), typegen_program) = try_join(
        || {
            let common_program = apply_common_transforms(
                project_name,
                Arc::clone(&program),
                Arc::clone(&connection_interface),
                Arc::clone(&base_fragment_names),
                Arc::clone(&perf_logger),
            )?;

            try_join(
                || {
                    let operation_program = apply_operation_transforms(
                        project_name,
                        Arc::clone(&common_program),
                        Arc::clone(&base_fragment_names),
                        Arc::clone(&perf_logger),
                    )?;

                    try_join(
                        || {
                            apply_normalization_transforms(
                                project_name,
                                Arc::clone(&operation_program),
                                Arc::clone(&perf_logger),
                            )
                        },
                        || {
                            apply_operation_text_transforms(
                                project_name,
                                Arc::clone(&operation_program),
                                Arc::clone(&perf_logger),
                            )
                        },
                    )
                },
                || {
                    apply_reader_transforms(
                        project_name,
                        Arc::clone(&common_program),
                        Arc::clone(&base_fragment_names),
                        Arc::clone(&perf_logger),
                    )
                },
            )
        },
        || {
            apply_typegen_transforms(
                project_name,
                Arc::clone(&program),
                Arc::clone(&base_fragment_names),
                Arc::clone(&perf_logger),
            )
        },
    )?;

    Ok(Programs {
        source: program,
        reader: reader_program,
        normalization: normalization_program,
        operation_text: text_program,
        typegen: typegen_program,
    })
}

/// Applies transforms that apply to every output.
fn apply_common_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    connection_interface: Arc<ConnectionInterface>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // * DisallowIdAsAlias (in validate)
    // + ConnectionTransform
    // * RelayDirectiveTransform (in validate)
    // + MaskTransform
    // + MatchTransform
    // + RefetchableFragmentTransform
    // + DeferStreamTransform
    let log_event = perf_logger.create_event("apply_common_transforms");
    log_event.string("project", project_name.to_string());
    let program = log_event.time("transform_connections", || {
        transform_connections(&program, connection_interface)
    });
    let program = log_event.time("mask", || mask(&program));
    let program = log_event.time("transform_match", || transform_match(&program))?;
    let program = log_event.time("transform_defer_stream", || {
        transform_defer_stream(&program)
    })?;
    let program = log_event.time("transform_refetchable_fragment", || {
        transform_refetchable_fragment(&program, &base_fragment_names, false)
    })?;
    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// Applies transforms only for generated reader code.
/// Corresponds to the "fragment transforms" in the JS compiler.
fn apply_reader_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + ClientExtensionsTransform
    // + FieldHandleTransform
    // + InlineDataFragmentTransform
    // + FlattenTransform, flattenAbstractTypes: true
    // + SkipRedundantNodesTransform
    let log_event = perf_logger.create_event("apply_reader_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("client_extensions", || client_extensions(&program));
    let program = log_event.time("handle_field_transform", || {
        handle_field_transform(&program)
    });
    let program = log_event.time("inline_data_fragment", || inline_data_fragment(&program))?;
    let program = log_event.time("remove_base_fragments", || {
        remove_base_fragments(&program, base_fragment_names)
    });
    let program = log_event.time("flatten", || flatten(&program, true));
    let program = log_event.time("skip_redundant_nodes", || skip_redundant_nodes(&program));

    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// Applies transforms that apply to all operation artifacts.
/// Corresponds to the "query transforms" in the JS compiler.
fn apply_operation_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + SplitModuleImportTransform
    // * ValidateUnusedVariablesTransform (Moved to common_transforms)
    // + ApplyFragmentArgumentTransform
    // - ValidateGlobalVariablesTransform
    // + GenerateIDFieldTransform
    // * TestOperationTransform - part of relay_codegen
    let log_event = perf_logger.create_event("apply_operation_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("split_module_import", || {
        split_module_import(&program, &base_fragment_names)
    });
    let program = log_event.time("apply_fragment_arguments", || {
        apply_fragment_arguments(&program)
    })?;
    let program = log_event.time("generate_id_field", || generate_id_field(&program));

    // TODO(T67052528): execute FB-specific transforms only if config options is provided
    let program = log_event.time("generate_preloadable_metadata", || {
        generate_preloadable_metadata(&program)
    });
    let program = log_event.time("generate_subscription_name_metadata", || {
        generate_subscription_name_metadata(&program)
    })?;
    let program = log_event.time("generate_live_query_metadata", || {
        generate_live_query_metadata(&program)
    })?;

    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// After the operation transforms, this applies further transforms that only
/// apply to the generated normalization code.
///
/// Corresponds to the "codegen transforms" in the JS compiler
fn apply_normalization_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + SkipUnreachableNodeTransform
    // + InlineFragmentsTransform
    // + ClientExtensionsTransform
    // + GenerateTypeNameTransform
    // + FlattenTransform, flattenAbstractTypes: true
    // + SkipRedundantNodesTransform
    let log_event = perf_logger.create_event("apply_normalization_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("relay_early_flush", || relay_early_flush(&program))?;
    let program = log_event.time("skip_unreachable_node", || skip_unreachable_node(&program));
    let program = log_event.time("inline_fragments", || inline_fragments(&program));
    let program = log_event.time("client_extensions", || client_extensions(&program));
    let program = log_event.time("generate_typename", || generate_typename(&program, true));
    let program = log_event.time("flatten", || flatten(&program, true));
    log_event.time("validate_module_conflicts", || {
        validate_module_conflicts(&program)
    })?;
    let program = log_event.time("skip_redundant_nodes", || skip_redundant_nodes(&program));
    let program = log_event.time("dedupe_type_discriminator", || {
        dedupe_type_discriminator(&program)
    });

    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// After the operation transforms, this applies further transforms that only
/// apply to the printed operation text.
///
/// Corresponds to the "print transforms" in the JS compiler
fn apply_operation_text_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + SkipSplitOperationTransform
    // - ClientExtensionsTransform
    // + SkipClientExtensionsTransform
    // + SkipUnreachableNodeTransform
    // + GenerateTypeNameTransform
    // + FlattenTransform, flattenAbstractTypes: false
    // - SkipHandleFieldTransform
    // - FilterDirectivesTransform
    // + SkipUnusedVariablesTransform
    // - ValidateRequiredArgumentsTransform
    let log_event = perf_logger.create_event("apply_operation_text_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("relay_early_flush", || relay_early_flush(&program))?;
    let program = log_event.time("skip_split_operation", || skip_split_operation(&program));
    let program = log_event.time("skip_client_extensions", || {
        skip_client_extensions(&program)
    });
    let program = log_event.time("skip_unreachable_node", || skip_unreachable_node(&program));
    let program = log_event.time("generate_typename", || generate_typename(&program, false));
    let program = log_event.time("flatten", || flatten(&program, false));
    let program = log_event.time("skip_unused_variables", || skip_unused_variables(&program));
    let program = log_event.time("unwrap_custom_directive_selection", || {
        unwrap_custom_directive_selection(&program)
    });
    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

fn apply_typegen_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // * RelayDirectiveTransform
    // + MaskTransform
    // + MatchTransform
    // + FlattenTransform, flattenAbstractTypes: false
    // + RefetchableFragmentTransform,
    let log_event = perf_logger.create_event("apply_typegen_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("mask", || mask(&program));
    let program = log_event.time("transform_match", || transform_match(&program))?;
    let program = log_event.time("flatten", || flatten(&program, false));
    let program = log_event.time("transform_refetchable_fragment", || {
        transform_refetchable_fragment(&program, &base_fragment_names, true)
            .expect("Expected errors to be validated in common transforms.")
    });
    let program = log_event.time("remove_base_fragments", || {
        remove_base_fragments(&program, base_fragment_names)
    });
    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

fn try_join<T1, F1, T2, F2, E>(f1: F1, f2: F2) -> Result<(T1, T2), Vec<E>>
where
    F1: FnOnce() -> Result<T1, Vec<E>> + Send,
    F2: FnOnce() -> Result<T2, Vec<E>> + Send,
    T1: Send,
    T2: Send,
    E: Send,
{
    let (v1, v2) = rayon::join(f1, f2);
    Ok((v1?, v2?))
}
<commit_msg>remove panic in typegen_transform<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{PerfLogEvent, PerfLogger};
use fnv::FnvHashSet;
use graphql_ir::{Program, ValidationResult};
use graphql_transforms::{
    apply_fragment_arguments, client_extensions, dedupe_type_discriminator, flatten,
    generate_id_field, generate_live_query_metadata, generate_preloadable_metadata,
    generate_subscription_name_metadata, generate_typename, handle_field_transform,
    inline_data_fragment, inline_fragments, mask, relay_early_flush, remove_base_fragments,
    skip_client_extensions, skip_redundant_nodes, skip_split_operation, skip_unreachable_node,
    skip_unused_variables, split_module_import, transform_connections, transform_defer_stream,
    transform_match, transform_refetchable_fragment, unwrap_custom_directive_selection,
    validate_module_conflicts, ConnectionInterface,
};
use interner::StringKey;
use std::sync::Arc;

pub struct Programs {
    pub source: Arc<Program>,
    pub reader: Arc<Program>,
    pub normalization: Arc<Program>,
    pub operation_text: Arc<Program>,
    pub typegen: Arc<Program>,
}

pub fn apply_transforms<TPerfLogger>(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    connection_interface: Arc<ConnectionInterface>,
    perf_logger: Arc<TPerfLogger>,
) -> ValidationResult<Programs>
where
    TPerfLogger: PerfLogger + 'static,
{
    // The execution pipeline is as follows, where items at the same indentation
    // can be computed independently and therefore in parallel:
    // |- common
    //  |- reader
    //  |- operation
    //     |- normalization
    //     |- operation_text
    // |- typegen
    //
    // NOTE: try_join(f1, f2) prefers the errors from f1 over f2, so process the normalization
    // program first since it is likely to include more errors (both per-fragment and
    // whole-operation errors) than the reader_program which just includes per-fragment errors.
    let (((normalization_program, text_program), reader_program), typegen_program) = try_join(
        || {
            let common_program = apply_common_transforms(
                project_name,
                Arc::clone(&program),
                Arc::clone(&connection_interface),
                Arc::clone(&base_fragment_names),
                Arc::clone(&perf_logger),
            )?;

            try_join(
                || {
                    let operation_program = apply_operation_transforms(
                        project_name,
                        Arc::clone(&common_program),
                        Arc::clone(&base_fragment_names),
                        Arc::clone(&perf_logger),
                    )?;

                    try_join(
                        || {
                            apply_normalization_transforms(
                                project_name,
                                Arc::clone(&operation_program),
                                Arc::clone(&perf_logger),
                            )
                        },
                        || {
                            apply_operation_text_transforms(
                                project_name,
                                Arc::clone(&operation_program),
                                Arc::clone(&perf_logger),
                            )
                        },
                    )
                },
                || {
                    apply_reader_transforms(
                        project_name,
                        Arc::clone(&common_program),
                        Arc::clone(&base_fragment_names),
                        Arc::clone(&perf_logger),
                    )
                },
            )
        },
        || {
            apply_typegen_transforms(
                project_name,
                Arc::clone(&program),
                Arc::clone(&base_fragment_names),
                Arc::clone(&perf_logger),
            )
        },
    )?;

    Ok(Programs {
        source: program,
        reader: reader_program,
        normalization: normalization_program,
        operation_text: text_program,
        typegen: typegen_program,
    })
}

/// Applies transforms that apply to every output.
fn apply_common_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    connection_interface: Arc<ConnectionInterface>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // * DisallowIdAsAlias (in validate)
    // + ConnectionTransform
    // * RelayDirectiveTransform (in validate)
    // + MaskTransform
    // + MatchTransform
    // + RefetchableFragmentTransform
    // + DeferStreamTransform
    let log_event = perf_logger.create_event("apply_common_transforms");
    log_event.string("project", project_name.to_string());
    let program = log_event.time("transform_connections", || {
        transform_connections(&program, connection_interface)
    });
    let program = log_event.time("mask", || mask(&program));
    let program = log_event.time("transform_match", || transform_match(&program))?;
    let program = log_event.time("transform_defer_stream", || {
        transform_defer_stream(&program)
    })?;
    let program = log_event.time("transform_refetchable_fragment", || {
        transform_refetchable_fragment(&program, &base_fragment_names, false)
    })?;
    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// Applies transforms only for generated reader code.
/// Corresponds to the "fragment transforms" in the JS compiler.
fn apply_reader_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + ClientExtensionsTransform
    // + FieldHandleTransform
    // + InlineDataFragmentTransform
    // + FlattenTransform, flattenAbstractTypes: true
    // + SkipRedundantNodesTransform
    let log_event = perf_logger.create_event("apply_reader_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("client_extensions", || client_extensions(&program));
    let program = log_event.time("handle_field_transform", || {
        handle_field_transform(&program)
    });
    let program = log_event.time("inline_data_fragment", || inline_data_fragment(&program))?;
    let program = log_event.time("remove_base_fragments", || {
        remove_base_fragments(&program, base_fragment_names)
    });
    let program = log_event.time("flatten", || flatten(&program, true));
    let program = log_event.time("skip_redundant_nodes", || skip_redundant_nodes(&program));

    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// Applies transforms that apply to all operation artifacts.
/// Corresponds to the "query transforms" in the JS compiler.
fn apply_operation_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + SplitModuleImportTransform
    // * ValidateUnusedVariablesTransform (Moved to common_transforms)
    // + ApplyFragmentArgumentTransform
    // - ValidateGlobalVariablesTransform
    // + GenerateIDFieldTransform
    // * TestOperationTransform - part of relay_codegen
    let log_event = perf_logger.create_event("apply_operation_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("split_module_import", || {
        split_module_import(&program, &base_fragment_names)
    });
    let program = log_event.time("apply_fragment_arguments", || {
        apply_fragment_arguments(&program)
    })?;
    let program = log_event.time("generate_id_field", || generate_id_field(&program));

    // TODO(T67052528): execute FB-specific transforms only if config options is provided
    let program = log_event.time("generate_preloadable_metadata", || {
        generate_preloadable_metadata(&program)
    });
    let program = log_event.time("generate_subscription_name_metadata", || {
        generate_subscription_name_metadata(&program)
    })?;
    let program = log_event.time("generate_live_query_metadata", || {
        generate_live_query_metadata(&program)
    })?;

    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// After the operation transforms, this applies further transforms that only
/// apply to the generated normalization code.
///
/// Corresponds to the "codegen transforms" in the JS compiler
fn apply_normalization_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + SkipUnreachableNodeTransform
    // + InlineFragmentsTransform
    // + ClientExtensionsTransform
    // + GenerateTypeNameTransform
    // + FlattenTransform, flattenAbstractTypes: true
    // + SkipRedundantNodesTransform
    let log_event = perf_logger.create_event("apply_normalization_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("relay_early_flush", || relay_early_flush(&program))?;
    let program = log_event.time("skip_unreachable_node", || skip_unreachable_node(&program));
    let program = log_event.time("inline_fragments", || inline_fragments(&program));
    let program = log_event.time("client_extensions", || client_extensions(&program));
    let program = log_event.time("generate_typename", || generate_typename(&program, true));
    let program = log_event.time("flatten", || flatten(&program, true));
    log_event.time("validate_module_conflicts", || {
        validate_module_conflicts(&program)
    })?;
    let program = log_event.time("skip_redundant_nodes", || skip_redundant_nodes(&program));
    let program = log_event.time("dedupe_type_discriminator", || {
        dedupe_type_discriminator(&program)
    });

    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

/// After the operation transforms, this applies further transforms that only
/// apply to the printed operation text.
///
/// Corresponds to the "print transforms" in the JS compiler
fn apply_operation_text_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // + SkipSplitOperationTransform
    // - ClientExtensionsTransform
    // + SkipClientExtensionsTransform
    // + SkipUnreachableNodeTransform
    // + GenerateTypeNameTransform
    // + FlattenTransform, flattenAbstractTypes: false
    // - SkipHandleFieldTransform
    // - FilterDirectivesTransform
    // + SkipUnusedVariablesTransform
    // - ValidateRequiredArgumentsTransform
    let log_event = perf_logger.create_event("apply_operation_text_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("relay_early_flush", || relay_early_flush(&program))?;
    let program = log_event.time("skip_split_operation", || skip_split_operation(&program));
    let program = log_event.time("skip_client_extensions", || {
        skip_client_extensions(&program)
    });
    let program = log_event.time("skip_unreachable_node", || skip_unreachable_node(&program));
    let program = log_event.time("generate_typename", || generate_typename(&program, false));
    let program = log_event.time("flatten", || flatten(&program, false));
    let program = log_event.time("skip_unused_variables", || skip_unused_variables(&program));
    let program = log_event.time("unwrap_custom_directive_selection", || {
        unwrap_custom_directive_selection(&program)
    });
    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

fn apply_typegen_transforms(
    project_name: StringKey,
    program: Arc<Program>,
    base_fragment_names: Arc<FnvHashSet<StringKey>>,
    perf_logger: Arc<impl PerfLogger>,
) -> ValidationResult<Arc<Program>> {
    // JS compiler
    // * RelayDirectiveTransform
    // + MaskTransform
    // + MatchTransform
    // + FlattenTransform, flattenAbstractTypes: false
    // + RefetchableFragmentTransform,
    let log_event = perf_logger.create_event("apply_typegen_transforms");
    log_event.string("project", project_name.to_string());

    let program = log_event.time("mask", || mask(&program));
    let program = log_event.time("transform_match", || transform_match(&program))?;
    let program = log_event.time("flatten", || flatten(&program, false));
    let program = log_event.time("transform_refetchable_fragment", || {
        transform_refetchable_fragment(&program, &base_fragment_names, true)
    })?;
    let program = log_event.time("remove_base_fragments", || {
        remove_base_fragments(&program, base_fragment_names)
    });
    perf_logger.complete_event(log_event);

    Ok(Arc::new(program))
}

fn try_join<T1, F1, T2, F2, E>(f1: F1, f2: F2) -> Result<(T1, T2), Vec<E>>
where
    F1: FnOnce() -> Result<T1, Vec<E>> + Send,
    F2: FnOnce() -> Result<T2, Vec<E>> + Send,
    T1: Send,
    T2: Send,
    E: Send,
{
    let (v1, v2) = rayon::join(f1, f2);
    Ok((v1?, v2?))
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::common;
use diem_types::transaction::{ArgumentABI, ScriptABI, TypeArgumentABI};
use heck::{CamelCase, ShoutySnakeCase};
use move_core_types::language_storage::TypeTag;
use serde_generate::{
    indent::{IndentConfig, IndentedWriter},
    python3, CodeGeneratorConfig,
};

use std::{
    collections::BTreeMap,
    io::{Result, Write},
    path::PathBuf,
};

/// Output transaction builders in Python for the given ABIs.
pub fn output(
    out: &mut dyn Write,
    serde_package_name: Option<String>,
    diem_package_name: Option<String>,
    abis: &[ScriptABI],
) -> Result<()> {
    let mut emitter = PythonEmitter {
        out: IndentedWriter::new(out, IndentConfig::Space(4)),
        serde_package_name,
        diem_package_name,
    };
    emitter.output_script_call_enum_with_imports(abis)?;
    emitter.output_additional_imports()?;

    emitter.output_encode_method()?;
    emitter.output_decode_method()?;

    for abi in abis {
        emitter.output_script_encoder_function(abi)?;
    }
    for abi in abis {
        emitter.output_script_decoder_function(abi)?;
    }

    for abi in abis {
        emitter.output_code_constant(abi)?;
    }
    emitter.output_encoder_map(abis)?;
    emitter.output_decoder_map(abis)?;

    emitter.output_decoding_helpers(abis)?;

    Ok(())
}

/// Shared state for the Python code generator.
struct PythonEmitter<T> {
    /// Writer.
    out: IndentedWriter<T>,
    /// Package where to find the serde module (if any).
    serde_package_name: Option<String>,
    /// Package where to find the diem module (if any).
    diem_package_name: Option<String>,
}

impl<T> PythonEmitter<T>
where
    T: Write,
{
    fn output_additional_imports(&mut self) -> Result<()> {
        writeln!(
            self.out,
            r#"
from {}diem_types import (Script, TypeTag, AccountAddress, TransactionArgument, TransactionArgument__Bool, TransactionArgument__U8, TransactionArgument__U64, TransactionArgument__U128, TransactionArgument__Address, TransactionArgument__U8Vector)"#,
            match &self.diem_package_name {
                None => "".into(),
                Some(package) => package.clone() + ".",
            },
        )
    }

    fn output_encode_method(&mut self) -> Result<()> {
        writeln!(
            self.out,
            r#"
def encode_script(call: ScriptCall) -> Script:
    """Build a Diem `Script` from a structured object `ScriptCall`.
    """
    helper = SCRIPT_ENCODER_MAP[call.__class__]
    return helper(call)
"#
        )
    }

    fn output_decode_method(&mut self) -> Result<()> {
        writeln!(
            self.out,
            r#"
def decode_script(script: Script) -> ScriptCall:
    """Try to recognize a Diem `Script` and convert it into a structured object `ScriptCall`.
    """
    helper = SCRIPT_DECODER_MAP.get(script.code)
    if helper is None:
        raise ValueError("Unknown script bytecode")
    return helper(script)
"#
        )
    }

    fn output_script_call_enum_with_imports(&mut self, abis: &[ScriptABI]) -> Result<()> {
        let diem_types_module = match &self.diem_package_name {
            None => "diem_types".into(),
            Some(package) => format!("{}.diem_types", package),
        };
        let external_definitions = crate::common::get_external_definitions(&diem_types_module);
        let script_registry: BTreeMap<_, _> = vec![(
            "ScriptCall".to_string(),
            crate::common::make_abi_enum_container(abis),
        )]
        .into_iter()
        .collect();
        let mut comments: BTreeMap<_, _> = abis
            .iter()
            .map(|abi| {
                (
                    vec![
                        "".to_string(),
                        "ScriptCall".to_string(),
                        abi.name().to_camel_case(),
                    ],
                    Self::prepare_doc_string(abi.doc()),
                )
            })
            .collect();
        comments.insert(
            vec!["".to_string(), "ScriptCall".to_string()],
            "Structured representation of a call into a known Move script.".into(),
        );
        let config = CodeGeneratorConfig::new("".to_string())
            .with_comments(comments)
            .with_external_definitions(external_definitions)
            .with_serialization(false);
        python3::CodeGenerator::new(&config)
            .with_serde_package_name(self.serde_package_name.clone())
            .output(&mut self.out, &script_registry)
            .map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, format!("{}", err)))?;
        Ok(())
    }

    fn output_script_encoder_function(&mut self, abi: &ScriptABI) -> Result<()> {
        writeln!(
            self.out,
            "\ndef encode_{}_script({}) -> Script:",
            abi.name(),
            [
                Self::quote_type_parameters(abi.ty_args()),
                Self::quote_parameters(abi.args()),
            ]
            .concat()
            .join(", ")
        )?;
        self.out.indent();
        writeln!(
            self.out,
            "\"\"\"{}\n\"\"\"",
            Self::prepare_doc_string(abi.doc())
        )?;
        writeln!(
            self.out,
            r#"return Script(
    code={}_CODE,
    ty_args=[{}],
    args=[{}],
)
"#,
            abi.name().to_shouty_snake_case(),
            Self::quote_type_arguments(abi.ty_args()),
            Self::quote_arguments(abi.args()),
        )?;
        self.out.unindent();
        Ok(())
    }

    fn output_script_decoder_function(&mut self, abi: &ScriptABI) -> Result<()> {
        writeln!(
            self.out,
            "\ndef decode_{}_script({}script: Script) -> ScriptCall:",
            abi.name(),
            // prevent warning "unused variable"
            if abi.ty_args().is_empty() && abi.args().is_empty() {
                "_"
            } else {
                ""
            }
        )?;
        self.out.indent();
        writeln!(
            self.out,
            "return ScriptCall__{0}(",
            abi.name().to_camel_case(),
        )?;
        self.out.indent();
        for (index, ty_arg) in abi.ty_args().iter().enumerate() {
            writeln!(self.out, "{}=script.ty_args[{}],", ty_arg.name(), index,)?;
        }
        for (index, arg) in abi.args().iter().enumerate() {
            writeln!(
                self.out,
                "{}=decode_{}_argument(script.args[{}]),",
                arg.name(),
                common::mangle_type(arg.type_tag()),
                index,
            )?;
        }
        self.out.unindent();
        writeln!(self.out, ")\n")?;
        self.out.unindent();
        Ok(())
    }

    fn output_code_constant(&mut self, abi: &ScriptABI) -> Result<()> {
        writeln!(
            self.out,
            "\n{}_CODE = b\"{}\"",
            abi.name().to_shouty_snake_case(),
            abi.code()
                .iter()
                .map(|x| format!("\\x{:02x}", x))
                .collect::<Vec<_>>()
                .join(""),
        )
    }

    fn output_encoder_map(&mut self, abis: &[ScriptABI]) -> Result<()> {
        writeln!(
            self.out,
            r#"
# pyre-ignore
SCRIPT_ENCODER_MAP: typing.Dict[typing.Type[ScriptCall], typing.Callable[[ScriptCall], Script]] = {{"#
        )?;
        self.out.indent();
        for abi in abis {
            writeln!(
                self.out,
                "ScriptCall__{}: encode_{}_script,",
                abi.name().to_camel_case(),
                abi.name()
            )?;
        }
        self.out.unindent();
        writeln!(self.out, "}}\n")
    }

    fn output_decoder_map(&mut self, abis: &[ScriptABI]) -> Result<()> {
        writeln!(
            self.out,
            "\nSCRIPT_DECODER_MAP: typing.Dict[bytes, typing.Callable[[Script], ScriptCall]] = {{"
        )?;
        self.out.indent();
        for abi in abis {
            writeln!(
                self.out,
                "{}_CODE: decode_{}_script,",
                abi.name().to_shouty_snake_case(),
                abi.name()
            )?;
        }
        self.out.unindent();
        writeln!(self.out, "}}\n")
    }

    fn output_decoding_helpers(&mut self, abis: &[ScriptABI]) -> Result<()> {
        let required_types = common::get_required_decoding_helper_types(abis);
        for required_type in required_types {
            self.output_decoding_helper(required_type)?;
        }
        Ok(())
    }

    fn output_decoding_helper(&mut self, type_tag: &TypeTag) -> Result<()> {
        use TypeTag::*;
        let (constructor, expr) = match type_tag {
            Bool => ("Bool", "arg.value".into()),
            U8 => ("U8", "arg.value".into()),
            U64 => ("U64", "arg.value".into()),
            U128 => ("U128", "arg.value".into()),
            Address => ("Address", "arg.value".into()),
            Vector(type_tag) => match type_tag.as_ref() {
                U8 => ("U8Vector", "arg.value".into()),
                inner_type_tag => (
                    "Vector",
                    format!(
                        "[decode_{}_argument(x) for x in arg.value]",
                        common::mangle_type(inner_type_tag)
                    ),
                ),
            },
            Struct(_) | Signer => common::type_not_allowed(type_tag),
        };
        writeln!(
            self.out,
            r#"
def decode_{}_argument(arg: TransactionArgument) -> {}:
    if not isinstance(arg, TransactionArgument__{}):
        raise ValueError("Was expecting a {} argument")
    return {}
"#,
            common::mangle_type(type_tag),
            Self::quote_type(type_tag),
            constructor,
            constructor,
            expr,
        )
    }

    fn prepare_doc_string(doc: &str) -> String {
        let doc = crate::common::prepare_doc_string(doc);
        let s: Vec<_> = doc.splitn(2, |c| c == '.').collect();
        if s.len() <= 1 || s[1].is_empty() {
            format!("{}.", s[0])
        } else {
            format!("{}.\n\n{}", s[0], s[1].trim())
        }
    }

    fn quote_type_parameters(ty_args: &[TypeArgumentABI]) -> Vec<String> {
        ty_args
            .iter()
            .map(|ty_arg| format!("{}: TypeTag", ty_arg.name()))
            .collect()
    }

    fn quote_parameters(args: &[ArgumentABI]) -> Vec<String> {
        args.iter()
            .map(|arg| format!("{}: {}", arg.name(), Self::quote_type(arg.type_tag())))
            .collect()
    }

    fn quote_type_arguments(ty_args: &[TypeArgumentABI]) -> String {
        ty_args
            .iter()
            .map(|ty_arg| ty_arg.name().to_string())
            .collect::<Vec<_>>()
            .join(", ")
    }

    fn quote_arguments(args: &[ArgumentABI]) -> String {
        args.iter()
            .map(|arg| Self::quote_transaction_argument(arg.type_tag(), arg.name()))
            .collect::<Vec<_>>()
            .join(", ")
    }

    fn quote_type(type_tag: &TypeTag) -> String {
        use TypeTag::*;
        match type_tag {
            Bool => "st.bool".into(),
            U8 => "st.uint8".into(),
            U64 => "st.uint64".into(),
            U128 => "st.uint128".into(),
            Address => "AccountAddress".into(),
            Vector(type_tag) => match type_tag.as_ref() {
                U8 => "bytes".into(),
                _ => common::type_not_allowed(type_tag),
            },

            Struct(_) | Signer => common::type_not_allowed(type_tag),
        }
    }

    fn quote_transaction_argument(type_tag: &TypeTag, name: &str) -> String {
        use TypeTag::*;
        match type_tag {
            Bool => format!("TransactionArgument__Bool(value={})", name),
            U8 => format!("TransactionArgument__U8(value={})", name),
            U64 => format!("TransactionArgument__U64(value={})", name),
            U128 => format!("TransactionArgument__U128(value={})", name),
            Address => format!("TransactionArgument__Address(value={})", name),
            Vector(type_tag) => match type_tag.as_ref() {
                U8 => format!("TransactionArgument__U8Vector(value={})", name),
                _ => common::type_not_allowed(type_tag),
            },

            Struct(_) | Signer => common::type_not_allowed(type_tag),
        }
    }
}

pub struct Installer {
    install_dir: PathBuf,
    serde_package_name: Option<String>,
    diem_package_name: Option<String>,
}

impl Installer {
    pub fn new(
        install_dir: PathBuf,
        serde_package_name: Option<String>,
        diem_package_name: Option<String>,
    ) -> Self {
        Installer {
            install_dir,
            serde_package_name,
            diem_package_name,
        }
    }

    fn open_module_init_file(&self, name: &str) -> Result<std::fs::File> {
        let dir_path = self.install_dir.join(name);
        std::fs::create_dir_all(&dir_path)?;
        std::fs::File::create(dir_path.join("__init__.py"))
    }
}

impl crate::SourceInstaller for Installer {
    type Error = Box<dyn std::error::Error>;

    fn install_transaction_builders(
        &self,
        name: &str,
        abis: &[ScriptABI],
    ) -> std::result::Result<(), Self::Error> {
        let mut file = self.open_module_init_file(name)?;
        output(
            &mut file,
            self.serde_package_name.clone(),
            self.diem_package_name.clone(),
            abis,
        )?;
        Ok(())
    }
}
<commit_msg>[transaction builder generator] fix new pyre error<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::common;
use diem_types::transaction::{ArgumentABI, ScriptABI, TypeArgumentABI};
use heck::{CamelCase, ShoutySnakeCase};
use move_core_types::language_storage::TypeTag;
use serde_generate::{
    indent::{IndentConfig, IndentedWriter},
    python3, CodeGeneratorConfig,
};

use std::{
    collections::BTreeMap,
    io::{Result, Write},
    path::PathBuf,
};

/// Output transaction builders in Python for the given ABIs.
pub fn output(
    out: &mut dyn Write,
    serde_package_name: Option<String>,
    diem_package_name: Option<String>,
    abis: &[ScriptABI],
) -> Result<()> {
    let mut emitter = PythonEmitter {
        out: IndentedWriter::new(out, IndentConfig::Space(4)),
        serde_package_name,
        diem_package_name,
    };
    emitter.output_script_call_enum_with_imports(abis)?;
    emitter.output_additional_imports()?;

    emitter.output_encode_method()?;
    emitter.output_decode_method()?;

    for abi in abis {
        emitter.output_script_encoder_function(abi)?;
    }
    for abi in abis {
        emitter.output_script_decoder_function(abi)?;
    }

    for abi in abis {
        emitter.output_code_constant(abi)?;
    }
    emitter.output_encoder_map(abis)?;
    emitter.output_decoder_map(abis)?;

    emitter.output_decoding_helpers(abis)?;

    Ok(())
}

/// Shared state for the Python code generator.
struct PythonEmitter<T> {
    /// Writer.
    out: IndentedWriter<T>,
    /// Package where to find the serde module (if any).
    serde_package_name: Option<String>,
    /// Package where to find the diem module (if any).
    diem_package_name: Option<String>,
}

impl<T> PythonEmitter<T>
where
    T: Write,
{
    fn output_additional_imports(&mut self) -> Result<()> {
        writeln!(
            self.out,
            r#"
from {}diem_types import (Script, TypeTag, AccountAddress, TransactionArgument, TransactionArgument__Bool, TransactionArgument__U8, TransactionArgument__U64, TransactionArgument__U128, TransactionArgument__Address, TransactionArgument__U8Vector)"#,
            match &self.diem_package_name {
                None => "".into(),
                Some(package) => package.clone() + ".",
            },
        )
    }

    fn output_encode_method(&mut self) -> Result<()> {
        writeln!(
            self.out,
            r#"
def encode_script(call: ScriptCall) -> Script:
    """Build a Diem `Script` from a structured object `ScriptCall`.
    """
    helper = SCRIPT_ENCODER_MAP[call.__class__]
    return helper(call)
"#
        )
    }

    fn output_decode_method(&mut self) -> Result<()> {
        writeln!(
            self.out,
            r#"
def decode_script(script: Script) -> ScriptCall:
    """Try to recognize a Diem `Script` and convert it into a structured object `ScriptCall`.
    """
    helper = SCRIPT_DECODER_MAP.get(script.code)
    if helper is None:
        raise ValueError("Unknown script bytecode")
    return helper(script)
"#
        )
    }

    fn output_script_call_enum_with_imports(&mut self, abis: &[ScriptABI]) -> Result<()> {
        let diem_types_module = match &self.diem_package_name {
            None => "diem_types".into(),
            Some(package) => format!("{}.diem_types", package),
        };
        let external_definitions = crate::common::get_external_definitions(&diem_types_module);
        let script_registry: BTreeMap<_, _> = vec![(
            "ScriptCall".to_string(),
            crate::common::make_abi_enum_container(abis),
        )]
        .into_iter()
        .collect();
        let mut comments: BTreeMap<_, _> = abis
            .iter()
            .map(|abi| {
                (
                    vec![
                        "".to_string(),
                        "ScriptCall".to_string(),
                        abi.name().to_camel_case(),
                    ],
                    Self::prepare_doc_string(abi.doc()),
                )
            })
            .collect();
        comments.insert(
            vec!["".to_string(), "ScriptCall".to_string()],
            "Structured representation of a call into a known Move script.".into(),
        );
        let config = CodeGeneratorConfig::new("".to_string())
            .with_comments(comments)
            .with_external_definitions(external_definitions)
            .with_serialization(false);
        python3::CodeGenerator::new(&config)
            .with_serde_package_name(self.serde_package_name.clone())
            .output(&mut self.out, &script_registry)
            .map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, format!("{}", err)))?;
        Ok(())
    }

    fn output_script_encoder_function(&mut self, abi: &ScriptABI) -> Result<()> {
        writeln!(
            self.out,
            "\ndef encode_{}_script({}) -> Script:",
            abi.name(),
            [
                Self::quote_type_parameters(abi.ty_args()),
                Self::quote_parameters(abi.args()),
            ]
            .concat()
            .join(", ")
        )?;
        self.out.indent();
        writeln!(
            self.out,
            "\"\"\"{}\n\"\"\"",
            Self::prepare_doc_string(abi.doc())
        )?;
        writeln!(
            self.out,
            r#"return Script(
    code={}_CODE,
    ty_args=[{}],
    args=[{}],
)
"#,
            abi.name().to_shouty_snake_case(),
            Self::quote_type_arguments(abi.ty_args()),
            Self::quote_arguments(abi.args()),
        )?;
        self.out.unindent();
        Ok(())
    }

    fn output_script_decoder_function(&mut self, abi: &ScriptABI) -> Result<()> {
        writeln!(
            self.out,
            "\ndef decode_{}_script({}script: Script) -> ScriptCall:",
            abi.name(),
            // prevent warning "unused variable"
            if abi.ty_args().is_empty() && abi.args().is_empty() {
                "_"
            } else {
                ""
            }
        )?;
        self.out.indent();
        writeln!(
            self.out,
            "return ScriptCall__{0}(",
            abi.name().to_camel_case(),
        )?;
        self.out.indent();
        for (index, ty_arg) in abi.ty_args().iter().enumerate() {
            writeln!(self.out, "{}=script.ty_args[{}],", ty_arg.name(), index,)?;
        }
        for (index, arg) in abi.args().iter().enumerate() {
            writeln!(
                self.out,
                "{}=decode_{}_argument(script.args[{}]),",
                arg.name(),
                common::mangle_type(arg.type_tag()),
                index,
            )?;
        }
        self.out.unindent();
        writeln!(self.out, ")\n")?;
        self.out.unindent();
        Ok(())
    }

    fn output_code_constant(&mut self, abi: &ScriptABI) -> Result<()> {
        writeln!(
            self.out,
            "\n{}_CODE = b\"{}\"",
            abi.name().to_shouty_snake_case(),
            abi.code()
                .iter()
                .map(|x| format!("\\x{:02x}", x))
                .collect::<Vec<_>>()
                .join(""),
        )
    }

    fn output_encoder_map(&mut self, abis: &[ScriptABI]) -> Result<()> {
        writeln!(
            self.out,
            r#"
# pyre-ignore
SCRIPT_ENCODER_MAP: typing.Dict[typing.Type[ScriptCall], typing.Callable[[ScriptCall], Script]] = {{"#
        )?;
        self.out.indent();
        for abi in abis {
            writeln!(
                self.out,
                "ScriptCall__{}: encode_{}_script,",
                abi.name().to_camel_case(),
                abi.name()
            )?;
        }
        self.out.unindent();
        writeln!(self.out, "}}\n")
    }

    fn output_decoder_map(&mut self, abis: &[ScriptABI]) -> Result<()> {
        writeln!(
            self.out,
            "\nSCRIPT_DECODER_MAP: typing.Dict[bytes, typing.Callable[[Script], ScriptCall]] = {{"
        )?;
        self.out.indent();
        for abi in abis {
            writeln!(
                self.out,
                "{}_CODE: decode_{}_script,",
                abi.name().to_shouty_snake_case(),
                abi.name()
            )?;
        }
        self.out.unindent();
        writeln!(self.out, "}}\n")
    }

    fn output_decoding_helpers(&mut self, abis: &[ScriptABI]) -> Result<()> {
        let required_types = common::get_required_decoding_helper_types(abis);
        for required_type in required_types {
            self.output_decoding_helper(required_type)?;
        }
        Ok(())
    }

    fn output_decoding_helper(&mut self, type_tag: &TypeTag) -> Result<()> {
        use TypeTag::*;
        let (constructor, expr) = match type_tag {
            Bool => ("Bool", "arg.value".into()),
            U8 => ("U8", "arg.value".into()),
            U64 => ("U64", "arg.value".into()),
            U128 => ("U128", "arg.value".into()),
            Address => ("Address", "arg.value".into()),
            Vector(type_tag) => match type_tag.as_ref() {
                U8 => ("U8Vector", "arg.value".into()),
                inner_type_tag => (
                    "Vector",
                    format!(
                        "[decode_{}_argument(x) for x in arg.value]",
                        common::mangle_type(inner_type_tag)
                    ),
                ),
            },
            Struct(_) | Signer => common::type_not_allowed(type_tag),
        };
        writeln!(
            self.out,
            r#"
def decode_{}_argument(arg: TransactionArgument) -> {}:
    if not isinstance(arg, TransactionArgument__{}):
        raise ValueError("Was expecting a {} argument")
    return {}
"#,
            common::mangle_type(type_tag),
            Self::quote_type(type_tag),
            constructor,
            constructor,
            expr,
        )
    }

    fn prepare_doc_string(doc: &str) -> String {
        let doc = crate::common::prepare_doc_string(doc);
        let s: Vec<_> = doc.splitn(2, |c| c == '.').collect();
        if s.len() <= 1 || s[1].is_empty() {
            format!("{}.", s[0])
        } else {
            format!("{}.\n\n{}", s[0], s[1].trim())
        }
    }

    fn quote_type_parameters(ty_args: &[TypeArgumentABI]) -> Vec<String> {
        ty_args
            .iter()
            .map(|ty_arg| format!("{}: TypeTag", ty_arg.name()))
            .collect()
    }

    fn quote_parameters(args: &[ArgumentABI]) -> Vec<String> {
        args.iter()
            .map(|arg| format!("{}: {}", arg.name(), Self::quote_type(arg.type_tag())))
            .collect()
    }

    fn quote_type_arguments(ty_args: &[TypeArgumentABI]) -> String {
        ty_args
            .iter()
            .map(|ty_arg| ty_arg.name().to_string())
            .collect::<Vec<_>>()
            .join(", ")
    }

    fn quote_arguments(args: &[ArgumentABI]) -> String {
        args.iter()
            .map(|arg| Self::quote_transaction_argument(arg.type_tag(), arg.name()))
            .collect::<Vec<_>>()
            .join(", ")
    }

    fn quote_type(type_tag: &TypeTag) -> String {
        use TypeTag::*;
        match type_tag {
            Bool => "bool".into(),
            U8 => "st.uint8".into(),
            U64 => "st.uint64".into(),
            U128 => "st.uint128".into(),
            Address => "AccountAddress".into(),
            Vector(type_tag) => match type_tag.as_ref() {
                U8 => "bytes".into(),
                _ => common::type_not_allowed(type_tag),
            },

            Struct(_) | Signer => common::type_not_allowed(type_tag),
        }
    }

    fn quote_transaction_argument(type_tag: &TypeTag, name: &str) -> String {
        use TypeTag::*;
        match type_tag {
            Bool => format!("TransactionArgument__Bool(value={})", name),
            U8 => format!("TransactionArgument__U8(value={})", name),
            U64 => format!("TransactionArgument__U64(value={})", name),
            U128 => format!("TransactionArgument__U128(value={})", name),
            Address => format!("TransactionArgument__Address(value={})", name),
            Vector(type_tag) => match type_tag.as_ref() {
                U8 => format!("TransactionArgument__U8Vector(value={})", name),
                _ => common::type_not_allowed(type_tag),
            },

            Struct(_) | Signer => common::type_not_allowed(type_tag),
        }
    }
}

pub struct Installer {
    install_dir: PathBuf,
    serde_package_name: Option<String>,
    diem_package_name: Option<String>,
}

impl Installer {
    pub fn new(
        install_dir: PathBuf,
        serde_package_name: Option<String>,
        diem_package_name: Option<String>,
    ) -> Self {
        Installer {
            install_dir,
            serde_package_name,
            diem_package_name,
        }
    }

    fn open_module_init_file(&self, name: &str) -> Result<std::fs::File> {
        let dir_path = self.install_dir.join(name);
        std::fs::create_dir_all(&dir_path)?;
        std::fs::File::create(dir_path.join("__init__.py"))
    }
}

impl crate::SourceInstaller for Installer {
    type Error = Box<dyn std::error::Error>;

    fn install_transaction_builders(
        &self,
        name: &str,
        abis: &[ScriptABI],
    ) -> std::result::Result<(), Self::Error> {
        let mut file = self.open_module_init_file(name)?;
        output(
            &mut file,
            self.serde_package_name.clone(),
            self.diem_package_name.clone(),
            abis,
        )?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::ich::{self, StableHashingContext};
use crate::traits::specialization_graph;
use crate::ty::fast_reject;
use crate::ty::fold::TypeFoldable;
use crate::ty::{Ty, TyCtxt};
use rustc_hir as hir;
use rustc_hir::def_id::{CrateNum, DefId};
use rustc_hir::definitions::DefPathHash;
use rustc_hir::HirId;

use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
use rustc_errors::ErrorReported;
use rustc_macros::HashStable;
use std::collections::BTreeMap;

/// A trait's definition with type information.
#[derive(HashStable)]
pub struct TraitDef {
    // We already have the def_path_hash below, no need to hash it twice
    #[stable_hasher(ignore)]
    pub def_id: DefId,

    pub unsafety: hir::Unsafety,

    /// If `true`, then this trait had the `#[rustc_paren_sugar]`
    /// attribute, indicating that it should be used with `Foo()`
    /// sugar. This is a temporary thing -- eventually any trait will
    /// be usable with the sugar (or without it).
    pub paren_sugar: bool,

    pub has_auto_impl: bool,

    /// If `true`, then this trait has the `#[marker]` attribute, indicating
    /// that all its associated items have defaults that cannot be overridden,
    /// and thus `impl`s of it are allowed to overlap.
    pub is_marker: bool,

    /// Used to determine whether the standard library is allowed to specialize
    /// on this trait.
    pub specialization_kind: TraitSpecializationKind,

    /// The ICH of this trait's DefPath, cached here so it doesn't have to be
    /// recomputed all the time.
    pub def_path_hash: DefPathHash,
}

/// Whether this trait is treated specially by the standard library
/// specialization lint.
#[derive(HashStable, PartialEq, Clone, Copy, TyEncodable, TyDecodable)]
pub enum TraitSpecializationKind {
    /// The default. Specializing on this trait is not allowed.
    None,
    /// Specializing on this trait is allowed because it doesn't have any
    /// methods. For example `Sized` or `FusedIterator`.
    /// Applies to traits with the `rustc_unsafe_specialization_marker`
    /// attribute.
    Marker,
    /// Specializing on this trait is allowed because all of the impls of this
    /// trait are "always applicable". Always applicable means that if
    /// `X<'x>: T<'y>` for any lifetimes, then `for<'a, 'b> X<'a>: T<'b>`.
    /// Applies to traits with the `rustc_specialization_trait` attribute.
    AlwaysApplicable,
}

#[derive(Default)]
pub struct TraitImpls {
    blanket_impls: Vec<DefId>,
    /// Impls indexed by their simplified self type, for fast lookup.
    non_blanket_impls: FxHashMap<fast_reject::SimplifiedType, Vec<DefId>>,
}

impl<'tcx> TraitDef {
    pub fn new(
        def_id: DefId,
        unsafety: hir::Unsafety,
        paren_sugar: bool,
        has_auto_impl: bool,
        is_marker: bool,
        specialization_kind: TraitSpecializationKind,
        def_path_hash: DefPathHash,
    ) -> TraitDef {
        TraitDef {
            def_id,
            unsafety,
            paren_sugar,
            has_auto_impl,
            is_marker,
            specialization_kind,
            def_path_hash,
        }
    }

    pub fn ancestors(
        &self,
        tcx: TyCtxt<'tcx>,
        of_impl: DefId,
    ) -> Result<specialization_graph::Ancestors<'tcx>, ErrorReported> {
        specialization_graph::ancestors(tcx, self.def_id, of_impl)
    }
}

impl<'tcx> TyCtxt<'tcx> {
    pub fn for_each_impl<F: FnMut(DefId)>(self, def_id: DefId, mut f: F) {
        let impls = self.trait_impls_of(def_id);

        for &impl_def_id in impls.blanket_impls.iter() {
            f(impl_def_id);
        }

        for v in impls.non_blanket_impls.values() {
            for &impl_def_id in v {
                f(impl_def_id);
            }
        }
    }

    /// Iterate over every impl that could possibly match the
    /// self type `self_ty`.
    pub fn for_each_relevant_impl<F: FnMut(DefId)>(
        self,
        def_id: DefId,
        self_ty: Ty<'tcx>,
        mut f: F,
    ) {
        let impls = self.trait_impls_of(def_id);

        for &impl_def_id in impls.blanket_impls.iter() {
            f(impl_def_id);
        }

        // simplify_type(.., false) basically replaces type parameters and
        // projections with infer-variables. This is, of course, done on
        // the impl trait-ref when it is instantiated, but not on the
        // predicate trait-ref which is passed here.
        //
        // for example, if we match `S: Copy` against an impl like
        // `impl<T:Copy> Copy for Option<T>`, we replace the type variable
        // in `Option<T>` with an infer variable, to `Option<_>` (this
        // doesn't actually change fast_reject output), but we don't
        // replace `S` with anything - this impl of course can't be
        // selected, and as there are hundreds of similar impls,
        // considering them would significantly harm performance.

        // This depends on the set of all impls for the trait. That is
        // unfortunate. When we get red-green recompilation, we would like
        // to have a way of knowing whether the set of relevant impls
        // changed. The most naive
        // way would be to compute the Vec of relevant impls and see whether
        // it differs between compilations. That shouldn't be too slow by
        // itself - we do quite a bit of work for each relevant impl anyway.
        //
        // If we want to be faster, we could have separate queries for
        // blanket and non-blanket impls, and compare them separately.
        //
        // I think we'll cross that bridge when we get to it.
        if let Some(simp) = fast_reject::simplify_type(self, self_ty, true) {
            if let Some(impls) = impls.non_blanket_impls.get(&simp) {
                for &impl_def_id in impls {
                    f(impl_def_id);
                }
            }
        } else {
            for &impl_def_id in impls.non_blanket_impls.values().flatten() {
                f(impl_def_id);
            }
        }
    }

    /// Applies function to every impl that could possibly match the self type `self_ty` and returns
    /// the first non-none value.
    pub fn find_map_relevant_impl<T, F: Fn(DefId) -> Option<T>>(
        self,
        def_id: DefId,
        self_ty: Ty<'tcx>,
        f: F,
    ) -> Option<T> {
        let impls = self.trait_impls_of(def_id);

        for &impl_def_id in impls.blanket_impls.iter() {
            if let result @ Some(_) = f(impl_def_id) {
                return result;
            }
        }

        if let Some(simp) = fast_reject::simplify_type(self, self_ty, true) {
            if let Some(impls) = impls.non_blanket_impls.get(&simp) {
                for &impl_def_id in impls {
                    if let result @ Some(_) = f(impl_def_id) {
                        return result;
                    }
                }
            }
        } else {
            for &impl_def_id in impls.non_blanket_impls.values().flatten() {
                if let result @ Some(_) = f(impl_def_id) {
                    return result;
                }
            }
        }

        None
    }

    /// Returns an iterator containing all impls
    pub fn all_impls(self, def_id: DefId) -> impl Iterator<Item = DefId> + 'tcx {
        let TraitImpls { blanket_impls, non_blanket_impls } = self.trait_impls_of(def_id);

        blanket_impls.iter().chain(non_blanket_impls.iter().map(|(_, v)| v).flatten()).cloned()
    }
}

// Query provider for `all_local_trait_impls`.
pub(super) fn all_local_trait_impls<'tcx>(
    tcx: TyCtxt<'tcx>,
    krate: CrateNum,
) -> &'tcx BTreeMap<DefId, Vec<HirId>> {
    &tcx.hir_crate(krate).trait_impls
}

// Query provider for `trait_impls_of`.
pub(super) fn trait_impls_of_provider(tcx: TyCtxt<'_>, trait_id: DefId) -> TraitImpls {
    let mut impls = TraitImpls::default();

    // Traits defined in the current crate can't have impls in upstream
    // crates, so we don't bother querying the cstore.
    if !trait_id.is_local() {
        for &cnum in tcx.crates().iter() {
            for &(impl_def_id, simplified_self_ty) in
                tcx.implementations_of_trait((cnum, trait_id)).iter()
            {
                if let Some(simplified_self_ty) = simplified_self_ty {
                    impls
                        .non_blanket_impls
                        .entry(simplified_self_ty)
                        .or_default()
                        .push(impl_def_id);
                } else {
                    impls.blanket_impls.push(impl_def_id);
                }
            }
        }
    }

    for &hir_id in tcx.hir().trait_impls(trait_id) {
        let impl_def_id = tcx.hir().local_def_id(hir_id).to_def_id();

        let impl_self_ty = tcx.type_of(impl_def_id);
        if impl_self_ty.references_error() {
            continue;
        }

        if let Some(simplified_self_ty) = fast_reject::simplify_type(tcx, impl_self_ty, false) {
            impls.non_blanket_impls.entry(simplified_self_ty).or_default().push(impl_def_id);
        } else {
            impls.blanket_impls.push(impl_def_id);
        }
    }

    impls
}

impl<'a> HashStable<StableHashingContext<'a>> for TraitImpls {
    fn hash_stable(&self, hcx: &mut StableHashingContext<'a>, hasher: &mut StableHasher) {
        let TraitImpls { ref blanket_impls, ref non_blanket_impls } = *self;

        ich::hash_stable_trait_impls(hcx, hasher, blanket_impls, non_blanket_impls);
    }
}
<commit_msg>Reimplement for_each_relevant_impl on top of find_map...<commit_after>use crate::ich::{self, StableHashingContext};
use crate::traits::specialization_graph;
use crate::ty::fast_reject;
use crate::ty::fold::TypeFoldable;
use crate::ty::{Ty, TyCtxt};
use rustc_hir as hir;
use rustc_hir::def_id::{CrateNum, DefId};
use rustc_hir::definitions::DefPathHash;
use rustc_hir::HirId;

use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
use rustc_errors::ErrorReported;
use rustc_macros::HashStable;
use std::collections::BTreeMap;

/// A trait's definition with type information.
#[derive(HashStable)]
pub struct TraitDef {
    // We already have the def_path_hash below, no need to hash it twice
    #[stable_hasher(ignore)]
    pub def_id: DefId,

    pub unsafety: hir::Unsafety,

    /// If `true`, then this trait had the `#[rustc_paren_sugar]`
    /// attribute, indicating that it should be used with `Foo()`
    /// sugar. This is a temporary thing -- eventually any trait will
    /// be usable with the sugar (or without it).
    pub paren_sugar: bool,

    pub has_auto_impl: bool,

    /// If `true`, then this trait has the `#[marker]` attribute, indicating
    /// that all its associated items have defaults that cannot be overridden,
    /// and thus `impl`s of it are allowed to overlap.
    pub is_marker: bool,

    /// Used to determine whether the standard library is allowed to specialize
    /// on this trait.
    pub specialization_kind: TraitSpecializationKind,

    /// The ICH of this trait's DefPath, cached here so it doesn't have to be
    /// recomputed all the time.
    pub def_path_hash: DefPathHash,
}

/// Whether this trait is treated specially by the standard library
/// specialization lint.
#[derive(HashStable, PartialEq, Clone, Copy, TyEncodable, TyDecodable)]
pub enum TraitSpecializationKind {
    /// The default. Specializing on this trait is not allowed.
    None,
    /// Specializing on this trait is allowed because it doesn't have any
    /// methods. For example `Sized` or `FusedIterator`.
    /// Applies to traits with the `rustc_unsafe_specialization_marker`
    /// attribute.
    Marker,
    /// Specializing on this trait is allowed because all of the impls of this
    /// trait are "always applicable". Always applicable means that if
    /// `X<'x>: T<'y>` for any lifetimes, then `for<'a, 'b> X<'a>: T<'b>`.
    /// Applies to traits with the `rustc_specialization_trait` attribute.
    AlwaysApplicable,
}

#[derive(Default)]
pub struct TraitImpls {
    blanket_impls: Vec<DefId>,
    /// Impls indexed by their simplified self type, for fast lookup.
    non_blanket_impls: FxHashMap<fast_reject::SimplifiedType, Vec<DefId>>,
}

impl<'tcx> TraitDef {
    pub fn new(
        def_id: DefId,
        unsafety: hir::Unsafety,
        paren_sugar: bool,
        has_auto_impl: bool,
        is_marker: bool,
        specialization_kind: TraitSpecializationKind,
        def_path_hash: DefPathHash,
    ) -> TraitDef {
        TraitDef {
            def_id,
            unsafety,
            paren_sugar,
            has_auto_impl,
            is_marker,
            specialization_kind,
            def_path_hash,
        }
    }

    pub fn ancestors(
        &self,
        tcx: TyCtxt<'tcx>,
        of_impl: DefId,
    ) -> Result<specialization_graph::Ancestors<'tcx>, ErrorReported> {
        specialization_graph::ancestors(tcx, self.def_id, of_impl)
    }
}

impl<'tcx> TyCtxt<'tcx> {
    pub fn for_each_impl<F: FnMut(DefId)>(self, def_id: DefId, mut f: F) {
        let impls = self.trait_impls_of(def_id);

        for &impl_def_id in impls.blanket_impls.iter() {
            f(impl_def_id);
        }

        for v in impls.non_blanket_impls.values() {
            for &impl_def_id in v {
                f(impl_def_id);
            }
        }
    }

    /// Iterate over every impl that could possibly match the
    /// self type `self_ty`.
    pub fn for_each_relevant_impl<F: FnMut(DefId)>(
        self,
        def_id: DefId,
        self_ty: Ty<'tcx>,
        mut f: F,
    ) {
        let _: Option<()> = self.find_map_relevant_impl(def_id, self_ty, |did| {
            f(did);
            None
        });
    }

    /// Applies function to every impl that could possibly match the self type `self_ty` and returns
    /// the first non-none value.
    pub fn find_map_relevant_impl<T, F: FnMut(DefId) -> Option<T>>(
        self,
        def_id: DefId,
        self_ty: Ty<'tcx>,
        mut f: F,
    ) -> Option<T> {
        let impls = self.trait_impls_of(def_id);

        for &impl_def_id in impls.blanket_impls.iter() {
            if let result @ Some(_) = f(impl_def_id) {
                return result;
            }
        }

        // simplify_type(.., false) basically replaces type parameters and
        // projections with infer-variables. This is, of course, done on
        // the impl trait-ref when it is instantiated, but not on the
        // predicate trait-ref which is passed here.
        //
        // for example, if we match `S: Copy` against an impl like
        // `impl<T:Copy> Copy for Option<T>`, we replace the type variable
        // in `Option<T>` with an infer variable, to `Option<_>` (this
        // doesn't actually change fast_reject output), but we don't
        // replace `S` with anything - this impl of course can't be
        // selected, and as there are hundreds of similar impls,
        // considering them would significantly harm performance.

        // This depends on the set of all impls for the trait. That is
        // unfortunate. When we get red-green recompilation, we would like
        // to have a way of knowing whether the set of relevant impls
        // changed. The most naive
        // way would be to compute the Vec of relevant impls and see whether
        // it differs between compilations. That shouldn't be too slow by
        // itself - we do quite a bit of work for each relevant impl anyway.
        //
        // If we want to be faster, we could have separate queries for
        // blanket and non-blanket impls, and compare them separately.
        //
        // I think we'll cross that bridge when we get to it.
        if let Some(simp) = fast_reject::simplify_type(self, self_ty, true) {
            if let Some(impls) = impls.non_blanket_impls.get(&simp) {
                for &impl_def_id in impls {
                    if let result @ Some(_) = f(impl_def_id) {
                        return result;
                    }
                }
            }
        } else {
            for &impl_def_id in impls.non_blanket_impls.values().flatten() {
                if let result @ Some(_) = f(impl_def_id) {
                    return result;
                }
            }
        }

        None
    }

    /// Returns an iterator containing all impls
    pub fn all_impls(self, def_id: DefId) -> impl Iterator<Item = DefId> + 'tcx {
        let TraitImpls { blanket_impls, non_blanket_impls } = self.trait_impls_of(def_id);

        blanket_impls.iter().chain(non_blanket_impls.iter().map(|(_, v)| v).flatten()).cloned()
    }
}

// Query provider for `all_local_trait_impls`.
pub(super) fn all_local_trait_impls<'tcx>(
    tcx: TyCtxt<'tcx>,
    krate: CrateNum,
) -> &'tcx BTreeMap<DefId, Vec<HirId>> {
    &tcx.hir_crate(krate).trait_impls
}

// Query provider for `trait_impls_of`.
pub(super) fn trait_impls_of_provider(tcx: TyCtxt<'_>, trait_id: DefId) -> TraitImpls {
    let mut impls = TraitImpls::default();

    // Traits defined in the current crate can't have impls in upstream
    // crates, so we don't bother querying the cstore.
    if !trait_id.is_local() {
        for &cnum in tcx.crates().iter() {
            for &(impl_def_id, simplified_self_ty) in
                tcx.implementations_of_trait((cnum, trait_id)).iter()
            {
                if let Some(simplified_self_ty) = simplified_self_ty {
                    impls
                        .non_blanket_impls
                        .entry(simplified_self_ty)
                        .or_default()
                        .push(impl_def_id);
                } else {
                    impls.blanket_impls.push(impl_def_id);
                }
            }
        }
    }

    for &hir_id in tcx.hir().trait_impls(trait_id) {
        let impl_def_id = tcx.hir().local_def_id(hir_id).to_def_id();

        let impl_self_ty = tcx.type_of(impl_def_id);
        if impl_self_ty.references_error() {
            continue;
        }

        if let Some(simplified_self_ty) = fast_reject::simplify_type(tcx, impl_self_ty, false) {
            impls.non_blanket_impls.entry(simplified_self_ty).or_default().push(impl_def_id);
        } else {
            impls.blanket_impls.push(impl_def_id);
        }
    }

    impls
}

impl<'a> HashStable<StableHashingContext<'a>> for TraitImpls {
    fn hash_stable(&self, hcx: &mut StableHashingContext<'a>, hasher: &mut StableHasher) {
        let TraitImpls { ref blanket_impls, ref non_blanket_impls } = *self;

        ich::hash_stable_trait_impls(hcx, hasher, blanket_impls, non_blanket_impls);
    }
}
<|endoftext|>