typhoon-context-macro 0.1.0-alpha.12

TODO
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

use {
    super::GeneratorResult,
    crate::{context::Context, visitor::ContextVisitor, StagedGenerator},
    quote::{format_ident, quote},
    syn::{Expr, Ident},
    typhoon_syn::constraints::{ConstraintInit, ConstraintInitIfNeeded, ConstraintToken},
};

#[derive(Default)]
struct TokenChecks {
    mint: Option<Ident>,
    owner: Option<Expr>,
    has_init: bool,
}

impl ContextVisitor for TokenChecks {
    fn visit_token(&mut self, constraint: &ConstraintToken) -> Result<(), syn::Error> {
        match constraint {
            ConstraintToken::Mint(ident) => self.mint = Some(ident.clone()),
            ConstraintToken::Owner(expr) => self.owner = Some(expr.clone()),
        }
        Ok(())
    }

    fn visit_init(&mut self, _constraint: &ConstraintInit) -> Result<(), syn::Error> {
        self.has_init = true;
        Ok(())
    }

    fn visit_init_if_needed(
        &mut self,
        _constraint: &ConstraintInitIfNeeded,
    ) -> Result<(), syn::Error> {
        self.has_init = true;
        Ok(())
    }
}

pub struct TokenAccountGenerator<'a>(&'a Context);

impl<'a> TokenAccountGenerator<'a> {
    pub fn new(context: &'a Context) -> Self {
        Self(context)
    }
}

impl StagedGenerator for TokenAccountGenerator<'_> {
    fn append(&mut self, result: &mut GeneratorResult) -> Result<(), syn::Error> {
        for account in &self.0.accounts {
            let mut checks = TokenChecks::default();
            checks.visit_account(account)?;

            if (checks.owner.is_some() || checks.mint.is_some()) && !checks.has_init {
                let mut check_token = Vec::with_capacity(2);
                let name = &account.name;
                let var_name = format_ident!("{}_state", name);

                if let Some(owner) = checks.owner {
                    check_token.push(quote! {
                        if #var_name.owner() != #owner.key() {
                            return Err(ErrorCode::TokenConstraintViolated.into());
                        }
                    });
                }

                if let Some(mint) = checks.mint {
                    check_token.push(quote! {
                        if #var_name.mint() != #mint.key() {
                            return Err(ErrorCode::TokenConstraintViolated.into());
                        }
                    });
                }

                result.inside.extend(check_token);
            }
        }
        Ok(())
    }
}