consortium-tee-macros-impl 0.2.0

// Copyright 2026 Ethan Wu
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

use heck::ToPascalCase;
use proc_macro2::{Span, TokenStream as TokenStream2};
use quote::quote;
use syn::{
    Ident, Token, braced,
    parse::{Parse, ParseStream},
    punctuated::Punctuated,
};

struct TeeServiceInput {
    commands: Vec<Ident>,
    context: Option<Ident>,
}

mod keywords {
    syn::custom_keyword!(context);
    syn::custom_keyword!(commands);
}

impl Parse for TeeServiceInput {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut commands: Option<Vec<Ident>> = None;
        let mut context: Option<Ident> = None;

        while !input.is_empty() {
            let lookahead = input.lookahead1();
            if lookahead.peek(keywords::context) {
                if context.is_some() {
                    return Err(syn::Error::new(input.span(), "Duplicate 'context' section"));
                }
                input.parse::<keywords::context>()?;
                let ctx_ident: Ident = input.parse()?;
                context = Some(ctx_ident);
            } else if lookahead.peek(keywords::commands) {
                if commands.is_some() {
                    return Err(syn::Error::new(
                        input.span(),
                        "Duplicate 'commands' section",
                    ));
                }
                input.parse::<keywords::commands>()?;
                let content;
                braced!(content in input);
                let idents = Punctuated::<Ident, Token![,]>::parse_terminated(&content)?;
                if idents
                    .iter()
                    .any(|ident| ident.to_string().eq_ignore_ascii_case("unknown"))
                {
                    return Err(syn::Error::new(
                        input.span(),
                        "Command name 'Unknown' is reserved for the default case",
                    ));
                }
                commands = Some(idents.into_iter().collect());
            } else {
                return Err(lookahead.error());
            }
        }

        Ok(TeeServiceInput {
            commands: commands.unwrap_or_default(),
            context,
        })
    }
}

/// Input:
///
/// ```no_run,ignore
/// tee_service! {
///   (optional) context ContextStruct
///   commands { function_a, function_b }
/// }
/// ```
///
/// Generates:
/// 1. `Command` enum with PascalCase variants and discriminants starting at 0, plus a
///    `#[default] Unknown` catch-all at the end.
/// 2. `invoke_command` (annotated `#[invoke_command]`) that dispatches to each
///    `function_*_dispatched` wrapper generated by `#[tee_command]`.
pub fn tee_service(input: TokenStream2) -> TokenStream2 {
    let TeeServiceInput { commands, context } = match syn::parse2::<TeeServiceInput>(input) {
        Ok(input) => input,
        Err(err) => return err.into_compile_error(),
    };

    let variant_idents: Vec<Ident> = commands
        .iter()
        .map(|cmd| Ident::new(&cmd.to_string().to_pascal_case(), cmd.span()))
        .collect();

    let discriminants: Vec<_> = (0u64..)
        .take(commands.len())
        .map(|n| syn::LitInt::new(&n.to_string(), Span::call_site()))
        .collect();

    let unknown_discriminant = syn::LitInt::new(&commands.len().to_string(), Span::call_site());

    let command_enum = quote! {
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
        #[repr(u32)]
        pub enum Command {
            #(#variant_idents = #discriminants,)*
            #[default]
            Unknown = #unknown_discriminant,
        }
    };

    let impl_primitive = quote! {
        impl From<u32> for Command {
            fn from(val: u32) -> Self {
                match val {
                    #(#discriminants => Command::#variant_idents,)*
                    _ => Command::Unknown,
                }
            }
        }

        impl From<Command> for u32 {
            fn from(cmd: Command) -> Self {
                cmd as u32
            }
        }
    };

    let dispatch_fns: Vec<Ident> = commands
        .iter()
        .map(|cmd| Ident::new(&format!("{}_dispatched", cmd), cmd.span()))
        .collect();

    let invoke_fn = if let Some(ctx_ident) = context {
        quote! {
            #[cfg(feature = "ta")]
            pub fn invoke_command(
                ctx: &mut #ctx_ident,
                cmd: u32,
                params: &mut ::optee_utee::Parameters,
            ) -> ::optee_utee::Result<()> {
                match Command::from(cmd) {
                    #(Command::#variant_idents => #dispatch_fns(ctx, params),)*
                    Command::Unknown => Err(::optee_utee::ErrorKind::BadParameters.into()),
                }
            }
        }
    } else {
        quote! {
            #[cfg(feature = "ta")]
            pub fn invoke_command(
                cmd: u32,
                params: &mut ::optee_utee::Parameters,
            ) -> ::optee_utee::Result<()> {
                match Command::from(cmd) {
                    #(Command::#variant_idents => #dispatch_fns(params),)*
                    Command::Unknown => Err(::optee_utee::ErrorKind::BadParameters.into()),
                }
            }
        }
    };

    quote! {
        #command_enum
        #impl_primitive
        #invoke_fn
    }
}