use crate::handle::{expr, item};
use proc_macro2::TokenStream;
use syn::parse::{Parse, ParseStream, Result};
use syn::{Expr, Ident, ItemFn, Token};
pub enum Scope {
Expr(Expr),
Item(ItemFn),
Verbatim(TokenStream), }
impl Scope {
pub fn handle(&self, args: &Ident) -> TokenStream {
match self {
Scope::Expr(scope) => expr::handle_expr(scope, args),
Scope::Item(scope) => item::handle_item_fn(scope, args),
Scope::Verbatim(_) => panic!("Custom functions not yet supported"),
}
}
}
impl Parse for Scope {
fn parse(input: ParseStream) -> Result<Self> {
if let Ok(item) = input.parse() {
Ok(Self::Item(item))
} else if let Ok(expr) = input.parse() {
Ok(Self::Expr(expr))
} else {
Ok(Scope::Verbatim(input.parse()?))
}
}
}
pub struct AutodiffAst {
pub args: Ident,
pub split: Token![:],
pub scope: Scope,
}
impl Parse for AutodiffAst {
fn parse(input: ParseStream) -> Result<Self> {
let args = input.parse()?;
let split = input.parse::<Token![:]>()?;
let scope = input.parse()?;
Ok(Self { args, split, scope })
}
}