// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#![recursion_limit = "128"]

extern crate proc_macro;

use proc_macro2::{Ident, TokenStream};
use quote::quote;
use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Index, Member, Variant};

#[cfg(test)]
mod tests;

// The method for packing an enum into a u64 is as follows:
// 1) Reserve the lowest "ceil(log_2(x))" bits where x is the number of enum variants.
// 2) Store the enum variant's index (0-based index based on order in the enum definition) in
//    reserved bits.
// 3) If there is data in the enum variant, store the data in remaining bits.
// The method for unpacking is as follows
// 1) Mask the raw token to just the reserved bits
// 2) Match the reserved bits to the enum variant token.
// 3) If the indicated enum variant had data, extract it from the unreserved bits.

// Calculates the number of bits needed to store the variant index. Essentially the log base 2
// of the number of variants, rounded up.
fn variant_bits(variants: &[Variant]) -> u32 {
    if variants.is_empty() {
        // The degenerate case of no variants.
        0
    } else {
        variants.len().next_power_of_two().trailing_zeros()
    }
}

// Name of the field if it has one, otherwise 0 assuming this is the zeroth
// field of a tuple variant.
fn field_member(field: &Field) -> Member {
    match &field.ident {
        Some(name) => Member::Named(name.clone()),
        None => Member::Unnamed(Index::from(0)),
    }
}

// Generates the function body for `as_raw_token`.
fn generate_as_raw_token(enum_name: &Ident, variants: &[Variant]) -> TokenStream {
    let variant_bits = variant_bits(variants);

    // Each iteration corresponds to one variant's match arm.
    let cases = variants.iter().enumerate().map(|(index, variant)| {
        let variant_name = &variant.ident;
        let index = index as u64;

        // The capture string is for everything between the variant identifier and the `=>` in
        // the match arm: the variant's data capture.
        let capture = variant.fields.iter().next().map(|field| {
            let member = field_member(&field);
            quote!({ #member: data })
        });

        // The modifier string ORs the variant index with extra bits from the variant data
        // field.
        let modifier = match variant.fields {
            Fields::Named(_) | Fields::Unnamed(_) => Some(quote! {
                | ((data as u64) << #variant_bits)
            }),
            Fields::Unit => None,
        };

        // Assembly of the match arm.
        quote! {
            #enum_name::#variant_name #capture => #index #modifier
        }
    });

    quote! {
        match *self {
            #(
                #cases,
            )*
        }
    }
}

// Generates the function body for `from_raw_token`.
fn generate_from_raw_token(enum_name: &Ident, variants: &[Variant]) -> TokenStream {
    let variant_bits = variant_bits(variants);
    let variant_mask = ((1 << variant_bits) - 1) as u64;

    // Each iteration corresponds to one variant's match arm.
    let cases = variants.iter().enumerate().map(|(index, variant)| {
        let variant_name = &variant.ident;
        let index = index as u64;

        // The data string is for extracting the enum variant's data bits out of the raw token
        // data, which includes both variant index and data bits.
        let data = variant.fields.iter().next().map(|field| {
            let member = field_member(&field);
            let ty = &field.ty;
            quote!({ #member: (data >> #variant_bits) as #ty })
        });

        // Assembly of the match arm.
        quote! {
            #index => #enum_name::#variant_name #data
        }
    });

    quote! {
        // The match expression only matches the bits for the variant index.
        match data & #variant_mask {
            #(
                #cases,
            )*
            _ => unreachable!(),
        }
    }
}

// The proc_macro::TokenStream type can only be constructed from within a
// procedural macro, meaning that unit tests are not able to invoke `fn
// poll_token` below as an ordinary Rust function. We factor out the logic into
// a signature that deals with Syn and proc-macro2 types only which are not
// restricted to a procedural macro invocation.
fn poll_token_inner(input: DeriveInput) -> TokenStream {
    let variants: Vec<Variant> = match input.data {
        Data::Enum(data) => data.variants.into_iter().collect(),
        Data::Struct(_) | Data::Union(_) => panic!("input must be an enum"),
    };

    for variant in &variants {
        assert!(variant.fields.iter().count() <= 1);
    }

    // Given our basic model of a user given enum that is suitable as a token, we generate the
    // implementation. The implementation is NOT always well formed, such as when a variant's data
    // type is not bit shiftable or castable to u64, but we let Rust generate such errors as it
    // would be difficult to detect every kind of error. Importantly, every implementation that we
    // generate here and goes on to compile succesfully is sound.

    let enum_name = input.ident;
    let as_raw_token = generate_as_raw_token(&enum_name, &variants);
    let from_raw_token = generate_from_raw_token(&enum_name, &variants);

    quote! {
        impl PollToken for #enum_name {
            fn as_raw_token(&self) -> u64 {
                #as_raw_token
            }

            fn from_raw_token(data: u64) -> Self {
                #from_raw_token
            }
        }
    }
}

/// Implements the PollToken trait for a given `enum`.
///
/// There are limitations on what `enum`s this custom derive will work on:
///
/// * Each variant must be a unit variant (no data), or have a single (un)named data field.
/// * If a variant has data, it must be a primitive type castable to and from a `u64`.
/// * If a variant data has size greater than or equal to a `u64`, its most significant bits must be
///   zero. The number of bits truncated is equal to the number of bits used to store the variant
///   index plus the number of bits above 64.
#[proc_macro_derive(PollToken)]
pub fn poll_token(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    poll_token_inner(input).into()
}