typhoon-context-macro 0.1.0-alpha.2

TODO
Documentation
use {
    super::{ConstraintGenerator, GeneratorResult},
    crate::{
        accounts::Account,
        constraints::{ConstraintAssociatedToken, ConstraintInit, ConstraintInitIfNeeded},
        visitor::ContextVisitor,
    },
    quote::quote,
};

#[derive(Default)]
pub struct RentGenerator {
    init_counter: i8,
    need_rent: bool,
}

impl RentGenerator {
    pub fn new() -> Self {
        Self::default()
    }
}

impl ConstraintGenerator for RentGenerator {
    fn generate(&self) -> Result<GeneratorResult, syn::Error> {
        if self.need_rent {
            Ok(GeneratorResult {
                at_init: quote! {
                    let rent = <Rent as Sysvar>::get()?;
                },
                ..Default::default()
            })
        } else {
            Ok(GeneratorResult::default())
        }
    }
}

impl ContextVisitor for RentGenerator {
    fn visit_account(&mut self, account: &Account) -> Result<(), syn::Error> {
        if !self.need_rent {
            self.init_counter = 0;
            self.visit_constraints(&account.constraints)?;
            if self.init_counter > 0 {
                self.need_rent = true;
            }
        }
        Ok(())
    }

    fn visit_init(&mut self, _contraint: &ConstraintInit) -> Result<(), syn::Error> {
        self.init_counter += 1;
        Ok(())
    }

    fn visit_init_if_needed(
        &mut self,
        _constraint: &ConstraintInitIfNeeded,
    ) -> Result<(), syn::Error> {
        self.init_counter += 1;

        Ok(())
    }

    fn visit_associated_token(
        &mut self,
        _constraint: &ConstraintAssociatedToken,
    ) -> Result<(), syn::Error> {
        self.init_counter -= 1;
        Ok(())
    }
}