// 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.

extern crate proc_macro;

use proc_macro::TokenStream;
use std::fmt::Write;
use std::mem;
use std::str::FromStr;

#[cfg(test)]
mod tests;

// This file is meant to be read from top to bottom to reflect how this code processes and generates
// Rust enum definitions and implementations. The algorithm overview:
// 1) Split the rust source by whitespace (`str::split_whitespace`).
// 2a) Attempt to tokenize each piece (see: `Tokenized`, all functions starting with `matches`).
// 2b) Feed the token to the `ParseState` (see `ParseState::handle_token`).
// 3) After the source is fully processed, the `ParseState` has an `EnumModel` representing the
//    input enum.
// 4) Glue together an implementation of PollToken using a template.

// A single token after splitting input source by white space and simple stateless matching.
#[derive(Clone, Debug, PartialEq)]
enum Tokenized {
    // `enum`
    Enum,

    // `pub' | `pub(crate)`
    Visiblity,

    // `Hello`, `index`, `data,`
    Ident(String),

    // `index:`, 'first:`
    FieldIdent(String),

    // `Socket(u32)`, `Client(usize),`,
    IdentAndType(String, String),

    // `{`
    OpenBrace,

    // `}`, `},`
    CloseBrace,
}

// Attempts to match strings of the form "identifier" with optional trailing comma.
fn matches_ident(s: &str) -> Option<String> {
    let ident = s.trim_right_matches(',');
    if !ident.is_empty() && ident.chars().all(char::is_alphanumeric) {
        Some(ident.to_owned())
    } else {
        None
    }
}

// Attempts to match strings of the form "Identifier(Type)" with optional trailing comma. If the
// given string matches, the identifier and type are returned as a 2-tuple receptively.
fn matches_ident_and_type(s: &str) -> Option<(String, String)> {
    let mut buffer = String::new();
    let mut ident = String::new();
    let mut type_ = String::new();
    let mut brace_depth = 0;
    for c in s.chars() {
        match c {
            '(' if brace_depth == 0 && !buffer.is_empty() && ident.is_empty() => {
                mem::swap(&mut ident, &mut buffer);
                brace_depth += 1;
            }
            ')' if brace_depth == 1 && !buffer.is_empty() && type_.is_empty() => {
                mem::swap(&mut type_, &mut buffer);
                brace_depth -= 1;
            }
            ',' => {}
            c if c.is_alphanumeric() => buffer.push(c),
            _ => return None,
        }
    }
    if !ident.is_empty() && !type_.is_empty() {
        Some((ident, type_))
    } else {
        None
    }
}

// Attempts to match strings of the form "identifier:".
fn matches_field_ident(s: &str) -> Option<String> {
    let field_ident = s.trim_right_matches(':');
    if s.ends_with(':') && field_ident.chars().all(char::is_alphanumeric) {
        Some(field_ident.to_owned())
    } else {
        None
    }
}

impl Tokenized {
    fn from_str(s: &str) -> Tokenized {
        if s.starts_with("pub(") {
            return Tokenized::Visiblity;
        }
        match s {
            "enum" => Tokenized::Enum,
            "pub" => Tokenized::Visiblity,
            "{" => Tokenized::OpenBrace,
            "}" | "}," => Tokenized::CloseBrace,
            _ => {
                // Try to match from most specific to least specific.
                if let Some(ident) = matches_field_ident(s) {
                    Tokenized::FieldIdent(ident)
                } else if let Some((ident, type_)) = matches_ident_and_type(s) {
                    Tokenized::IdentAndType(ident, type_)
                } else if let Some(ident) = matches_ident(s) {
                    Tokenized::Ident(ident)
                } else {
                    panic!("unable to parse token: {}", s)
                }
            }
        }
    }
}

// Data field for an enum, with possible field name.
#[derive(Debug, PartialEq)]
struct EnumVariantData {
    type_: String,
    name: Option<String>,
}

// Data for one variant of an enum, with optional single data field.
#[derive(Debug, PartialEq)]
struct EnumVariant {
    name: String,
    data: Option<EnumVariantData>,
}

// Data for an entire enum type.
#[derive(Debug, Default, PartialEq)]
struct EnumModel {
    name: String,
    variants: Vec<EnumVariant>,
}
// Note: impl for EnumModel is below the parsing code and definitions because all of the methods are
// for generating the PollToken impl.

// Labels for each of the states in the parsing state machine. The '->` symbol means that the given
// state may transition to the state pointed to.
#[derive(PartialEq, Debug)]
enum States {
    // Initial state, expecting to see visibility rules (e.g. `pub`) or `enum` keyword.
    Start, // -> Ident

    // Expect to see the name of the enum field.
    Ident, // -> Brace

    // Expect to see an opening brace.
    Brace, // -> VariantIdent, -> End

    // Expect to see a variant's name.
    VariantIdent, // -> VariantIdent, -> VariantData, -> End

    // Expect to see the field name of a variant's data.
    VariantData, // -> VariantIdent, -> VariantDataType

    // Expect to see the tye name of a variant's data.
    VariantDataType, // -> VariantData

    // Expect to see no more tokens.
    End,
}

// The state machine for parsing a stream of `Tokenized`. After the States::End state is reached, a
// complete `EnumModel` is ready to be used for generating an implementation.
struct ParseState {
    current_state: States,
    current_variant: Option<EnumVariant>,
    model: EnumModel,
}

impl ParseState {
    fn new() -> ParseState {
        ParseState {
            current_state: States::Start,
            current_variant: Default::default(),
            model: Default::default(),
        }
    }

    // Handles the next token in the stream of tokens.
    fn handle_token(&mut self, tok: Tokenized) {
        match self.current_state {
            States::Start => self.handle_start(tok),
            States::Ident => self.handle_ident(tok),
            States::Brace => self.handle_brace(tok),
            States::VariantIdent => self.handle_variant_ident(tok),
            States::VariantData => self.handle_variant_data(tok),
            States::VariantDataType => self.handle_variant_data_type(tok),
            States::End => self.handle_end(tok),
        }
    }

    // All the following are handlers name after the current state that handle the next token.

    fn handle_start(&mut self, tok: Tokenized) {
        self.current_state = match tok {
            Tokenized::Enum => States::Ident,
            Tokenized::Visiblity => States::Start,
            _ => panic!("derives for enum types only"),
        };
    }

    fn handle_ident(&mut self, tok: Tokenized) {
        self.current_state = match tok {
            Tokenized::Ident(ident) => {
                self.model.name = ident;
                States::Brace
            }
            _ => panic!("unexpected token: {:?}", tok),
        };
    }

    fn handle_brace(&mut self, tok: Tokenized) {
        self.current_state = match tok {
            Tokenized::OpenBrace => States::VariantIdent,
            Tokenized::CloseBrace => States::End,
            _ => panic!("unexpected token: {:?}", tok),
        };
    }

    fn handle_variant_ident(&mut self, tok: Tokenized) {
        // This handler is the most complex because it has the most branches for the new
        // `current_state`. Adding to that complexity is that many branches indicate a new variant
        // is being handled, which means the old `current_variant` needs to be added to `variants`
        // and a fresh one needs to be started with the fresh data embedded in the token.
        self.current_state = match tok {
            Tokenized::Ident(ident) => {
                let mut variant = Some(EnumVariant {
                    name: ident,
                    data: None,
                });
                mem::swap(&mut variant, &mut self.current_variant);
                if let Some(variant) = variant {
                    self.model.variants.push(variant);
                }
                States::VariantIdent
            }
            Tokenized::IdentAndType(ident, type_) => {
                let variant_data = EnumVariantData { type_, name: None };
                let mut variant = Some(EnumVariant {
                    name: ident,
                    data: Some(variant_data),
                });
                mem::swap(&mut variant, &mut self.current_variant);
                if let Some(variant) = variant {
                    self.model.variants.push(variant);
                }
                States::VariantIdent
            }
            Tokenized::OpenBrace => States::VariantData,
            Tokenized::CloseBrace => {
                let mut variant = Default::default();
                mem::swap(&mut variant, &mut self.current_variant);
                if let Some(variant) = variant {
                    self.model.variants.push(variant);
                }
                States::End
            }
            _ => panic!("unexpected token: {:?}", tok),
        };
    }

    fn handle_variant_data(&mut self, tok: Tokenized) {
        let variant = self.current_variant.as_mut().unwrap();
        self.current_state = match tok {
            Tokenized::FieldIdent(ident) => {
                assert!(
                    variant.data.is_none(),
                    "enum variant can only have one field"
                );
                variant.data = Some(EnumVariantData {
                    type_: "".to_owned(),
                    name: Some(ident),
                });
                States::VariantDataType
            }
            Tokenized::CloseBrace => States::VariantIdent,
            _ => panic!("unexpected token: {:?}", tok),
        };
    }

    fn handle_variant_data_type(&mut self, tok: Tokenized) {
        let variant = self.current_variant.as_mut().unwrap();
        let variant_data = variant.data.as_mut().unwrap();
        self.current_state = match tok {
            Tokenized::Ident(ident) => {
                variant_data.type_ = ident;
                States::VariantData
            }
            _ => panic!("unexpected token: {:?}", tok),
        };
    }

    fn handle_end(&mut self, tok: Tokenized) {
        panic!("unexpected tokens past ending brace: {:?}", tok);
    }
}

// Continued from the above `EnumModel` definition. All methods are used for generating PollToken
// implementation. 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.
impl EnumModel {
    // 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(&self) -> u32 {
        // The degenerate case of no variants.
        if self.variants.is_empty() {
            return 0;
        }
        self.variants.len().next_power_of_two().trailing_zeros()
    }

    // Generates the function body for `as_raw_token`.
    fn generate_as_raw_token(&self) -> String {
        let variant_bits = self.variant_bits();
        let mut match_statement = "match *self {\n".to_owned();

        // Each iteration corresponds to one variant's match arm.
        for (index, variant) in self.variants.iter().enumerate() {
            // The capture string is for everything between the variant identifier and the `=>` in
            // the match arm: the variant's data capture.
            let capture = match variant.data.as_ref() {
                Some(&EnumVariantData {
                    name: Some(ref name),
                    ..
                }) => format!("{{ {}: data }}", name),
                Some(&EnumVariantData { .. }) => "(data)".to_owned(),
                None => "".to_owned(),
            };

            // The modifier string ORs the variant index with extra bits from the variant data
            // field.
            let modifer = if variant.data.is_some() {
                format!(" | ((data as u64) << {})", variant_bits)
            } else {
                "".to_owned()
            };

            // Assembly of the match arm.
            write!(
                match_statement,
                "{}::{}{} => {}{},\n",
                self.name, variant.name, capture, index, modifer
            ).unwrap();
        }
        match_statement.push_str("}");
        match_statement
    }

    // Generates the function body for `from_raw_token`.
    fn generate_from_raw_token(&self) -> String {
        let variant_bits = self.variant_bits();
        let variant_mask = (1 << variant_bits) - 1;

        // The match expression only matches the bits for the variant index.
        let mut match_statement = format!("match data & 0x{:02x} {{\n", variant_mask);

        // Each iteration corresponds to one variant's match arm.
        for (index, variant) in self.variants.iter().enumerate() {
            // 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 = match variant.data.as_ref() {
                Some(&EnumVariantData {
                    name: Some(ref name),
                    ref type_,
                }) => format!("{{ {}: (data >> {}) as {} }}", name, variant_bits, type_),
                Some(&EnumVariantData {
                    name: None,
                    ref type_,
                }) => format!("((data >> {}) as {})", variant_bits, type_),
                None => "".to_owned(),
            };

            // Assembly of the match arm.
            write!(
                match_statement,
                "{} => {}::{}{},\n",
                index, self.name, variant.name, data
            ).unwrap();
        }
        match_statement.push_str("_ => unreachable!()\n}");
        match_statement
    }
}

// Because unit tests cannot create `TokenStream`s (apparently), we have an inner implementation
// that deals in strings.
fn poll_token_inner(src: &str) -> String {
    let src_tokens = src.split_whitespace();

    // Parsing is done in two interleaved stages, tokenizing without context, followed by parsing
    // via state machine.
    let mut state = ParseState::new();
    for src_tok in src_tokens {
        let tok = Tokenized::from_str(src_tok);
        state.handle_token(tok);
    }

    assert_eq!(
        state.current_state,
        States::End,
        "unexpected end after parsing source enum"
    );

    // 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 model = state.model;
    format!(
        "impl PollToken for {} {{
    fn as_raw_token(&self) -> u64 {{
{}
    }}

    fn from_raw_token(data: u64) -> Self {{
{}
    }}
}}",
        model.name,
        model.generate_as_raw_token(),
        model.generate_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: TokenStream) -> TokenStream {
    // The token stream gets converted to a string in a rather regular way, which makes parsing
    // simpler. In particular, whitespace from the source enum is not carried over, instead replaced
    // with whatever the token stream's to_string function outputs. The rust parser has already
    // validated the syntax, so we can make lots of assumptions about the source being well formed.
    TokenStream::from_str(&poll_token_inner(&input.to_string())).unwrap()
}