// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

#[derive(Debug, thiserror::Error)]
pub enum ParserError {
    #[error("failed to set tree_sitter language: {0}")]
    Language(#[from] tree_sitter::LanguageError),
    #[error("parser timed out or was cancelled")]
    DidNotComplete,
}

#[derive(Debug, thiserror::Error)]
pub enum RecoverableError {
    #[error("encountered unexpected expression")]
    UnexpectedExpression(String),
    #[error("encountered unimplemented expression")]
    UnimplementedStatement(String),
    #[error("expected a node, but it was missing")]
    MissingChild, // TODO: add String parameter to state which child is missing
    #[error("expected comparison to have lhs node, but it was missing")]
    MissingLhs,
    #[error("expected BinaryOperator node, but got unexpected node kind: {0}")]
    MissingOperator(String),
    #[error("Syntax error resulted in ERROR node")]
    SyntaxError(String),
}

pub fn recoverable_error_to_string(recoverable_error: &RecoverableError) -> String {
    match recoverable_error {
        RecoverableError::UnexpectedExpression(expression_name) => {
            format!("UnexpectedExpression: {}", expression_name)
        }
        RecoverableError::UnimplementedStatement(statement_name) => {
            format!("UnimplementedStatement: {}", statement_name)
        }
        RecoverableError::MissingChild => "MissingChild".to_string(),
        RecoverableError::MissingLhs => "MissingLhs".to_string(),
        RecoverableError::MissingOperator(operator) => {
            format!("MissingOperator: {}", operator)
        }
        RecoverableError::SyntaxError(invalid_syntax) => {
            format!("SyntaxError: {}", invalid_syntax)
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

#[macro_use]
extern crate rust_to_ocaml_attr;

use cst_to_ast::Parser as CSTToASTParser;
use parser_pre_process::remove_comments;

use crate::errors::recoverable_error_to_string;

pub mod ast;
pub mod constants;
pub mod cst_to_ast;
pub mod errors;
pub mod node_wrapper;
pub mod sitter;
pub mod string_helpers;

ocamlrep_ocamlpool::ocaml_ffi! {
    fn parse_module(_source_text: bstr::BString) -> Result< (ast::Mod_, Vec<ast::RecoverableErrorWithLocation>), String> {
        let input_code_as_rust_string = format!("{}", _source_text);
        let input_without_comments = remove_comments(input_code_as_rust_string);

        let mut cst_to_ast = CSTToASTParser::new(input_without_comments);
        match cst_to_ast.parse() {
            Ok(_) => {
                let ast_and_metadata = cst_to_ast.ast_and_metadata;
                let ast = ast_and_metadata.ast.unwrap();
                let recoverable_errors = ast_and_metadata.recoverable_errors;
                let mut errors: Vec<ast::RecoverableErrorWithLocation> = vec![];
                for recoverable_error_with_location in recoverable_errors{
                    let error_message = recoverable_error_to_string(&recoverable_error_with_location.parser_error);
                    let location = &recoverable_error_with_location.location;

                    errors.push(ast::RecoverableErrorWithLocation{
                        error: error_message,
                        lineno: location.lineno,
                        col_offset: location.col_offset,
                        end_lineno: location.end_lineno,
                        end_col_offset: location.end_col_offset
                    });
                }
                Ok((ast, errors))
            }
            Err(e) => Err(format!("Failure parsing input\n{:?}\n", e))
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

use pyo3::prelude::*;

#[macro_use]
extern crate rust_to_ocaml_attr;

use cst_to_ast::Parser as CSTToASTParser;

use crate::printers::parse_module_print_ast_code;
use crate::printers::parse_module_print_ast_pretty_and_errors;
use crate::printers::PrintingMode;

pub mod ast;
pub mod constants;
pub mod cst_to_ast;
pub mod errors;
pub mod node_wrapper;
pub mod parser_pre_process;
pub mod printers;
pub mod sitter;
pub mod string_helpers;

#[pyfunction]
fn py_parse_module_print_ast(input_code: String) -> PyResult<(String, String)> {
    Ok(parse_module_print_ast_code(
        input_code,
        PrintingMode::ASTOnly,
    ))
}

#[pyfunction]
fn py_parse_module_print_ast_and_pretty_print(input_code: String) -> PyResult<(String, String)> {
    Ok(parse_module_print_ast_code(
        input_code,
        PrintingMode::ASTAndPrettyPrintAST,
    ))
}

#[pyfunction]
fn py_parse_module_print_ast_pretty_print_only(input_code: String) -> PyResult<(String, String)> {
    Ok(parse_module_print_ast_code(
        input_code,
        PrintingMode::PrettyPrintASTOnly,
    ))
}

#[pyfunction]
fn py_parse_module_print_ast_pretty_and_errors(
    input_code: String,
) -> PyResult<(String, String, String)> {
    Ok(parse_module_print_ast_pretty_and_errors(input_code))
}

///
/// run errpy but ignore result - useful for benchmarking
///
#[pyfunction]
fn py_parse_module(input_code: String) -> PyResult<()> {
    let mut cst_to_ast = CSTToASTParser::new(input_code);
    _ = cst_to_ast.parse();

    Ok(())
}

/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn ffi_python(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(py_parse_module_print_ast, m)?)?;
    m.add_function(wrap_pyfunction!(
        py_parse_module_print_ast_and_pretty_print,
        m
    )?)?;
    m.add_function(wrap_pyfunction!(
        py_parse_module_print_ast_pretty_print_only,
        m
    )?)?;
    m.add_function(wrap_pyfunction!(
        py_parse_module_print_ast_pretty_and_errors,
        m
    )?)?;
    m.add_function(wrap_pyfunction!(py_parse_module, m)?)?;
    Ok(())
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

use std::collections::HashMap;
use std::collections::VecDeque;
use std::fmt;
use std::iter;
use std::str;

use tree_sitter::Node as TSNode;
use tree_sitter::Point;
use tree_sitter::TreeCursor;

/// Tree-sitter can emit `comment` and `error` nodes anywhere in its output CST's.
/// This means that in theory our code needs to handle the possibility of a comment
/// or error node being present anywhere in the CST! This is not really practical
/// to implement in `cst_to_ast.rs` so instead here we add `node_wrapper` to filter
/// out ERROR and COMMENT notes in a new `Node` struct which wrapps around
/// the Tree sitter `TSNode`.
///
/// We make use of an arena based tree allocator within `FilteredCST` for ease of
/// understandig and future maintenance.
///
/// The methods exposed here are a subset of those on the `TSNode` API visible at:
/// https://docs.rs/tree-sitter/latest/src/tree_sitter/lib.rs.html

#[derive(Debug, Default)]
pub struct FilteredCST<'tree> {
    nodes: Vec<Node<'tree>>,
}

pub struct Node<'a> {
    ts_node: TSNode<'a>,

    parent_idx: Option<usize>,
    children: Vec<usize>,
    named_children: Vec<usize>,
    children_for_field_name: HashMap<&'a str, Vec<usize>>,
}

impl<'a> fmt::Debug for Node<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self.ts_node)
    }
}

impl<'tree> FilteredCST<'tree> {
    /// Place node into FilteredCST tree and return numerial index
    fn allocate(&mut self, to_push: Node<'tree>) -> usize {
        self.nodes.push(to_push);
        self.nodes.len() - 1
    }

    pub fn get_root(&self) -> &Node<'tree> {
        &self.nodes[0]
    }
}

impl<'a> PartialEq for Node<'a> {
    fn eq(&self, other: &Self) -> bool {
        self.ts_node.id() == other.ts_node.id()
    }
}

/// iterative method to create a node tree from root TS Node
pub fn build_node_tree(root_ts_node: TSNode) -> FilteredCST {
    let mut arena = FilteredCST { nodes: vec![] };

    let mut to_process = VecDeque::new();
    to_process.push_back((None, None, root_ts_node));

    while !to_process.is_empty() {
        if let Some((parent_position, field_name, ts_node)) = to_process.pop_front() {
            let mut current_node = Node::new(ts_node);
            current_node.parent_idx = parent_position;
            let current_node_idx = arena.allocate(current_node);

            if let Some(parent_position) = parent_position {
                arena.nodes[parent_position].children.push(current_node_idx);
                if ts_node.is_named() {
                    arena.nodes[parent_position]
                        .named_children
                        .push(current_node_idx);
                }

                if let Some(field_name) = field_name {
                    let children_for_field_name =
                        &mut arena.nodes[parent_position].children_for_field_name;

                    if !children_for_field_name.contains_key(field_name) {
                        children_for_field_name.insert(field_name, vec![]);
                    }
                    children_for_field_name
                        .get_mut(&field_name)
                        .unwrap()
                        .push(current_node_idx);
                }
            }

            let field_names: Vec<Option<&str>> = extract_child_field_names(ts_node);

            let walker = &mut ts_node.walk();
            for (child, child_field_name) in ts_node.children(walker).zip(field_names.iter()) {
                let kind = child.kind();
                if kind == "ERROR" || kind == "comment" {
                    continue;
                }

                to_process.push_back((Some(current_node_idx), *child_field_name, child));
            }
        }
    }

    arena
}

/// Extract all child field names from a TSNode
fn extract_child_field_names<'a>(ts_node: TSNode<'a>) -> Vec<Option<&str>> {
    let cursor = &mut ts_node.walk();
    cursor.reset(ts_node);
    cursor.goto_first_child();
    (0..ts_node.child_count())
        .map(move |_| {
            let result = cursor.field_name();
            cursor.goto_next_sibling();
            result
        })
        .collect()
}

impl<'tree> Node<'tree> {
    fn new(ts_node: TSNode<'tree>) -> Self {
        Self {
            ts_node,

            parent_idx: None,
            children: vec![],
            named_children: vec![],
            children_for_field_name: HashMap::new(),
        }
    }

    /// Get a numeric id for this node that is unique.
    pub fn id(&self) -> usize {
        self.ts_node.id()
    }

    /// Get this node's start position in terms of rows and columns.
    pub fn start_position(&self) -> Point {
        self.ts_node.start_position()
    }

    /// Get this node's end position in terms of rows and columns.
    pub fn end_position(&self) -> Point {
        self.ts_node.end_position()
    }

    /// Get the node's child at the given index, where zero represents the first child.
    pub fn child<'a>(&'a self, filtered_cst: &'a FilteredCST<'tree>, idx: usize) -> Option<&Self> {
        self.children
            .get(idx)
            .map(|child_index| &filtered_cst.nodes[*child_index])
    }

    /// Get this node's number of children.
    pub fn child_count(&self) -> usize {
        self.children.len()
    }

    /// Get this node's utf8 text as a string.
    pub fn utf8_text<'a>(&self, source: &'a [u8]) -> Result<&'a str, str::Utf8Error> {
        self.ts_node.utf8_text(source)
    }

    /// Get the byte offsets where this node starts.
    pub fn start_byte(&self) -> usize {
        self.ts_node.start_byte()
    }

    /// Get the byte offsets where this node end.
    pub fn end_byte(&self) -> usize {
        self.ts_node.end_byte()
    }

    /// Create a new [TreeCursor] starting from this node.
    pub fn walk(&self) -> TreeCursor<'tree> {
        self.ts_node.walk()
    }

    /// Get this node's type as a string.
    pub fn kind(&self) -> &'static str {
        self.ts_node.kind()
    }

    /// Get this node's immediate parent.
    pub fn parent<'a>(&'a self, filtered_cst: &'a FilteredCST<'tree>) -> Option<&Self> {
        self.parent_idx
            .map(|parent_idx| &filtered_cst.nodes[parent_idx])
    }

    /// Check if this node is *named*.
    ///
    /// Named nodes correspond to named rules in the grammar, whereas *anonymous* nodes
    /// correspond to string literals in the grammar.
    pub fn is_named(&self) -> bool {
        self.ts_node.is_named()
    }

    /// Get the first child with the given field name.
    ///
    /// If multiple children may have the same field name, access them using
    /// [children_by_field_name](Node::children_by_field_name)
    pub fn child_by_field_name<'a>(
        &'a self,
        filtered_cst: &'a FilteredCST<'tree>,
        field_name: &'a str,
    ) -> Option<&'a Node<'a>> {
        self.children_by_field_name(filtered_cst, field_name).next()
    }

    /// Iterate over this node's children with a given field name.
    ///
    /// See also [Node::children].
    pub fn children_by_field_name<'a>(
        &'a self,
        filtered_cst: &'a FilteredCST<'tree>,
        field_name: &str,
    ) -> Box<dyn Iterator<Item = &Node<'a>> + 'a> {
        match self.children_for_field_name.get(field_name) {
            Some(children) => Box::new(
                children
                    .iter()
                    .map(|child_index| &filtered_cst.nodes[*child_index]),
            ),
            None => Box::new(iter::empty()),
        }
    }

    /// Iterate over this node's named children.
    ///
    /// See also [Node::children].
    pub fn named_children<'a>(
        &'a self,
        filtered_cst: &'a FilteredCST<'tree>,
    ) -> impl ExactSizeIterator<Item = &Node<'a>> + 'a {
        self.named_children
            .iter()
            .map(|child_index| &filtered_cst.nodes[*child_index])
    }

    /// Iterate over this node's children.
    pub fn children<'a>(
        &'a self,
        filtered_cst: &'a FilteredCST<'tree>,
    ) -> impl ExactSizeIterator<Item = &Node<'tree>> + 'a {
        self.children
            .iter()
            .map(|child_index| &filtered_cst.nodes[*child_index])
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

#[derive(PartialEq)]
enum StringCommentState {
    OutsideString,
    InsideSingleQuotedString,
    InsideDoubleQuotedString,
    InsideSingleQuotedStringEscape,
    InsideDoubleQuotedStringEscape,
    InsideSingleQuotedMultilineString,
    InsideDoubleQuotedMultilineString,
    InsideComment,
}

// handles a single (') or double quote (") and returns new parser state
fn handle_quote_character(
    input_code: &str,
    state: StringCommentState,
    idx: &usize,
    multiline_quote: &str,
    inside_string: StringCommentState,
    inside_string_escaped: StringCommentState,
    inside_multiline_string: StringCommentState,
) -> StringCommentState {
    use StringCommentState::*;

    if state == OutsideString {
        if input_code[..idx + 1].ends_with(multiline_quote) {
            // triple quote => start new multiline string
            inside_multiline_string
        } else {
            // single quote => we are now inside a string
            inside_string
        }
    } else if state == inside_string_escaped {
        // escaped quote => we are still inside the string
        inside_string
    } else if state == inside_string {
        // quote ends string => we are now outside
        OutsideString
    } else if state == inside_multiline_string {
        if input_code[..idx + 1].ends_with(multiline_quote)
            && (input_code.len() == idx + 1 || !input_code[..idx + 2].ends_with(multiline_quote))
        {
            // triple quote ends the multiline string => we are now outside
            OutsideString
        } else {
            // a single quote won't do anything in a multiline string
            // => we are still inside the multiline string
            inside_multiline_string
        }
    } else {
        // any other case (e.g. single quote in double quote string)
        // => don't change the parser state
        state
    }
}

///
/// This function will remove comment's from input_code.
/// Comments are replaced with newlines so as to preserve line
/// and column information
///
/// TODO: This code will need refactoring in the future if/when we wish to
///  * Handle `# pyre-fixme:`'s better by integrating them into the AST.
///  * Outut a CST code will need refactoring.
///
/// For now this function is only used on in the context of AST production
/// so it's fit for purpose.
///
pub fn remove_comments(input_code: String) -> String {
    use StringCommentState::*;

    let without_comments = &mut String::new();

    // start of current slice to copy to result string
    let mut start = 0;
    // line end index (either ends at comment start or at newline)
    let mut line_end: Option<usize> = None;
    let mut state = OutsideString;
    for (idx, char) in input_code.char_indices() {
        match char {
            '\\' => match state {
                InsideSingleQuotedString => state = InsideSingleQuotedStringEscape,
                InsideDoubleQuotedString => state = InsideDoubleQuotedStringEscape,
                _ => {}
            },
            '\'' => {
                state = handle_quote_character(
                    &input_code,
                    state,
                    &idx,
                    "'''",
                    InsideSingleQuotedString,
                    InsideSingleQuotedStringEscape,
                    InsideSingleQuotedMultilineString,
                )
            }
            '\"' => {
                state = handle_quote_character(
                    &input_code,
                    state,
                    &idx,
                    "\"\"\"",
                    InsideDoubleQuotedString,
                    InsideDoubleQuotedStringEscape,
                    InsideDoubleQuotedMultilineString,
                )
            }
            '#' => match state {
                OutsideString => {
                    // we are not inside a string and are at the start of a comment
                    // => line ends already _before_ the hash character
                    line_end = Some(idx);
                    state = InsideComment;
                }
                _ => {}
            },
            '\n' => {
                if state == InsideSingleQuotedStringEscape
                    || state == InsideDoubleQuotedStringEscape
                {
                    // newline was escaped within string
                    continue;
                }
                if line_end.is_none() {
                    // no line comment => end line here (at newline)
                    line_end = Some(idx);
                }
                // note that this slice does not contain the newline character as
                // the newline character will be contained in the next slice (`start` points to the newline char)
                without_comments.push_str(&input_code[start..line_end.unwrap()]);
                start = idx;
                line_end = None;
                match state {
                    InsideSingleQuotedMultilineString | InsideDoubleQuotedMultilineString => {}
                    _ => state = OutsideString,
                }
            }
            _ => {
                if state == InsideSingleQuotedStringEscape {
                    state = InsideSingleQuotedString;
                } else if state == InsideDoubleQuotedStringEscape {
                    state = InsideDoubleQuotedString;
                }
            }
        }
    }
    if line_end.is_none() {
        // no line comment at the end
        // => copy all remaining code into the result
        line_end = Some(input_code.len());
    }
    without_comments.push_str(&input_code[start..line_end.unwrap()]);
    // this extra newline is actually not needed, but kept for backwards-compatibility
    without_comments.push('\n');
    without_comments.to_string()
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

pub mod ast_pretty_print;
pub mod ast_pretty_print_helper;
pub mod ast_print;

use ast_pretty_print_helper::PrintHelper;
use cst_to_ast::get_node_text;
use cst_to_ast::ASTAndMetaData;
use cst_to_ast::Parser as CSTToASTParser;
use errors::recoverable_error_to_string;
use errors::ParserError;
use node_wrapper::build_node_tree;
use node_wrapper::Node;
use parser_pre_process::remove_comments;
use tree_sitter::Parser as TreeSitterParser;

use crate::cst_to_ast;
use crate::errors;
use crate::node_wrapper;
use crate::node_wrapper::FilteredCST;
use crate::parser_pre_process;

pub enum PrintingMode {
    ASTOnly,
    PrettyPrintASTOnly,
    ASTAndPrettyPrintAST,
}

pub fn parse_module_print_ast_pretty_and_errors(input_code: String) -> (String, String, String) {
    let parsing_and_printer = ParsingAndPrinter::new(input_code);
    (
        parsing_and_printer.get_print_ast(),
        parsing_and_printer.get_pretty_print_ast(),
        parsing_and_printer.get_recoverable_errors(),
    )
}

pub fn parse_module_print_ast_code(
    input_code: String,
    printing_mode: PrintingMode,
) -> (String, String) {
    let (print_ast, pretty_print_ast) = match printing_mode {
        PrintingMode::ASTOnly => (true, false),
        PrintingMode::PrettyPrintASTOnly => (false, true),
        PrintingMode::ASTAndPrettyPrintAST => (true, true),
    };

    let parsing_and_printer = ParsingAndPrinter::new(input_code);

    let pprint_output = &mut String::new();
    if print_ast {
        pprint_output.push_str(parsing_and_printer.get_print_ast().as_str());
    }
    if pretty_print_ast {
        pprint_output.push_str(parsing_and_printer.get_pretty_print_ast().as_str());
    }

    (
        pprint_output.to_string(),
        parsing_and_printer.get_recoverable_errors(),
    )
}

struct ParsingAndPrinter {
    cst_to_ast: CSTToASTParser,
    parse_result: Result<(), ParserError>,
}

impl ParsingAndPrinter {
    pub fn new(input_code: String) -> Self {
        let input_without_comments = remove_comments(input_code);
        let mut cst_to_ast = CSTToASTParser::new(input_without_comments);
        let parse_result = cst_to_ast.parse();
        ParsingAndPrinter {
            cst_to_ast,
            parse_result,
        }
    }

    pub fn get_print_ast(&self) -> String {
        match &self.parse_result {
            Ok(_) => format!("{}\n", self.cst_to_ast.get_ast_and_metadata()),
            Err(e) => format!("Failure parsing python module\n{:?}\n", e),
        }
    }

    pub fn get_pretty_print_ast(&self) -> String {
        let to_print = &mut String::new();
        let printer = &mut PrintHelper::new(to_print, 4);
        self.cst_to_ast.get_ast_and_metadata().pprint(printer);

        to_print.to_string()
    }

    pub fn get_recoverable_errors(&self) -> String {
        extract_errors(self.cst_to_ast.get_ast_and_metadata())
    }
}

///
/// Will return an empty string if there are no errors
///
fn extract_errors(metadata: &ASTAndMetaData) -> String {
    let mut errors = String::new();
    if !metadata.recoverable_errors.is_empty() {
        let error_count = metadata.recoverable_errors.len();

        let errors_word = if error_count > 1 { "errors" } else { "error" };

        errors.push_str(&format!(
            "{} Recoverable {} detected:\n",
            error_count, errors_word
        ));

        for recoverable_error_with_location in &metadata.recoverable_errors {
            let recoverable_as_string =
                recoverable_error_to_string(&recoverable_error_with_location.parser_error);

            let location = &recoverable_error_with_location.location;
            errors.push_str(
                &format!(
                    "{} at [{}:{} - {}:{}]",
                    recoverable_as_string,
                    location.lineno,
                    location.col_offset,
                    location.end_lineno,
                    location.end_col_offset
                )
                .to_string(),
            );
            errors.push('\n');

            let stack = &recoverable_error_with_location.stack;
            errors.push_str("CST Stack:\n");
            for item in stack.iter() {
                errors.push_str("    ");
                errors.push_str(item);
                errors.push('\n');
            }
        }
    }
    errors
}

pub struct CSTPrinter {
    code: String,
}

impl CSTPrinter {
    pub fn new(code: String) -> CSTPrinter {
        CSTPrinter { code }
    }

    // Prints the sitter nodes in their "derived `Debug` form"
    pub fn print_cst(&self) {
        let mut parser = TreeSitterParser::new();
        parser
            .set_language(tree_sitter_python::language())
            .expect("Fail to initialize TreeSitter");
        let tree = parser.parse(&self.code, None).expect("Fail to parse file");
        let filtered_cst = build_node_tree(tree.root_node());
        println!(">>> Tree-Sitter CST Nodes:\n");
        self.print_cst_node(&filtered_cst, filtered_cst.get_root(), "");
    }

    fn print_cst_node(&self, filtered_cst: &FilteredCST, node: &Node, indent: &str) {
        println!(
            "{}{:?} :: {}",
            indent,
            node,
            get_node_text(&self.code, node).replace('\n', "\\n")
        );
        for child in node.children(filtered_cst) {
            let new_indent = format!("  {}", indent);
            self.print_cst_node(filtered_cst, child, new_indent.as_str());
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

#[macro_use]
extern crate rust_to_ocaml_attr;

use crate::printers::parse_module_print_ast_code;
use crate::printers::PrintingMode;

pub mod ast;
pub mod constants;
pub mod cst_to_ast;
pub mod errors;
pub mod node_wrapper;
pub mod parser_pre_process;
pub mod printers;
pub mod sitter;
pub mod string_helpers;

/// Python Parser which will output AST and optional pretty print of input code from derived AST
#[derive(clap::Parser)]
struct Args {
    /// Python code to generate AST for input string
    input_code: String,
}

fn main() {
    let args = <Args as clap::Parser>::parse();

    let input_code = args.input_code;
    let (ast, errors) = parse_module_print_ast_code(input_code, PrintingMode::ASTAndPrettyPrintAST);
    if errors.is_empty() {
        println!("{}", ast);
    } else {
        println!("{}{}", ast, errors);
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

#[macro_use]
extern crate rust_to_ocaml_attr;

use clap::Parser as ClapParser;

use crate::printers::CSTPrinter;

pub mod ast;
pub mod constants;
pub mod cst_to_ast;
pub mod errors;
pub mod node_wrapper;
pub mod parser_pre_process;
pub mod printers;
pub mod sitter;
pub mod string_helpers;

/// Python Parser which will output Tree-sitter CST
#[derive(ClapParser)]
struct Args {
    /// Python code to generate CST for
    input_code: String,
}

fn main() {
    let args = <Args as clap::Parser>::parse();
    let input_code = args.input_code;
    CSTPrinter::new(input_code).print_cst();
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

#![allow(dead_code)]

/*!
 The structs below are used to wrap the CST nodes output by tree-sitter
 into something that is a bit easier to reason about...

 Tree sitter seems to have a string representation of tokens with
 numerical ids that are known at runtime...

 What we do here is to convert a sitter node to an internal
 representation and use that to `match` in Rust for later processing
 into an AST by the cst_to_ast.rs lowerer.

 This wrapping could be eliminated at some point or made cheaper.
*/

use crate::node_wrapper::Node;

///
/// Token supertype after extraction from Tree-sitter CST
///
#[derive(Debug)]
pub enum NodeType<'a> {
    Error,
    Keyword(Keyword),
    BinaryOperator(BinaryOperator),
    AugAssignOperator(AugAssignOperator),
    ComparisonOperator(ComparisonOperator),
    Delimiter(Delimiter),
    Production(Production<'a>),
}

#[derive(Copy, Clone, Debug)]
#[repr(u8)]
#[allow(non_camel_case_types)]
pub enum Keyword {
    AND,
    AS,
    ASSERT,
    ASYNC,
    AWAIT,
    BREAK,
    CONTINUE,
    CLASS,
    DEF,
    ELIF,
    ELSE,
    EXCEPT,
    DEL,
    FALSE,
    FINALLY,
    FOR,
    FROM,
    GLOBAL,
    IF,
    IMPORT,
    IN,
    IS,
    IS_NOT,
    LAMBDA,
    NONE,
    NONLOCAL,
    NOT,
    NOT_IN,
    OR,
    PASS,
    RAISE,
    RETURN,
    TRUE,
    TRY,
    YIELD,
    WHILE,
    WITH,
}

#[derive(Copy, Clone, Debug)]
#[repr(u8)]
#[allow(non_camel_case_types)]
pub enum BinaryOperator {
    ASSIGN_EXPR,
    AT,
    BITWISE_AND,
    BITWISE_OR,
    BITWISE_XOR,
    DOUBLE_SLASH,
    DOUBLE_STAR,
    LEFT_SHIFT,
    MINUS,
    PERCENT,
    PLUS,
    RIGHT_SHIFT,
    SLASH,
    STAR,
    TILDE,
}

#[derive(Copy, Clone, Debug)]
#[repr(u8)]
#[allow(non_camel_case_types)]
pub enum AugAssignOperator {
    AT_EQUAL,
    BITWISE_AND_EQUAL,
    BITWISE_OR_EQUAL,
    DOUBLE_SLASH_EQUAL,
    DOUBLE_STAR_EQUAL,
    LEFT_SHIFT_EQUAL,
    MINUS_EQUAL,
    PERCENT_EQUAL,
    PLUS_EQUAL,
    RIGHT_SHIFT_EQUAL,
    SLASH_EQUAL,
    STAR_EQUAL,
    BITWISE_XOR_EQUAL,
}

#[derive(Copy, Clone, Debug)]
#[repr(u8)]
#[allow(non_camel_case_types)]
pub enum ComparisonOperator {
    LESS_THAN,
    LESS_THAN_EQUAL,
    EQUAL,
    NOT_EQUAL,
    GREATER_THAN_EQUAL,
    GREATER_THAN,
    IN,
    NOT,
    IS,
}

#[derive(Copy, Clone, Debug)]
#[repr(u8)]
#[allow(non_camel_case_types)]
pub enum Delimiter {
    ARROW,
    ASSIGN,
    BACK_SLASH,
    COLON,
    COMMA,
    DUBLE_QUOTE,
    HASH,
    LEFT_BRACKET,
    LEFT_CURLY,
    LEFT_PAREN,
    PERIOD,
    RIGHT_BRACKET,
    RIGHT_CURLY,
    RIGHT_PAREN,
    SEMICOLON,
    SINGLE_QUOTE,
}

///
/// A production is a node to traverse and it carries a
/// reference to the sitter node
///
#[derive(Debug)]
pub struct Production<'a> {
    pub production_kind: ProductionKind,
    pub node: &'a Node<'a>,
}

///
/// Various forms of ProductionKind
///
#[derive(Debug, PartialEq)]
#[allow(non_camel_case_types)]
pub enum ProductionKind {
    // statements
    FUTURE_IMPORT_STATEMENT,
    IMPORT_STATEMENT,
    IMPORT_FROM_STATEMENT,
    PRINT_STATEMENT,
    ASSERT_STATEMENT,
    EXPRESSION_STATEMENT,
    RETURN_STATEMENT,
    DELETE_STATEMENT,
    RAISE_STATEMENT,
    PASS_STATEMENT,
    BREAK_STATEMENT,
    CONTINUE_STATEMENT,
    GLOBAL_STATEMENT,
    NONLOCAL_STATEMENT,
    EXEC_STATEMENT,
    IF_STATEMENT,
    FOR_STATEMENT,
    WHILE_STATEMENT,
    TRY_STATEMENT,
    WITH_STATEMENT,
    FUNCTION_DEFINITION,
    CLASS_DEFINITION,
    DECORATED_DEFINITION,
    MATCH_STATEMENT,

    // `expression_statement`
    ASSIGNMENT,
    AUGMENTED_ASSIGNMENT,
    YIELD,

    // `expression`
    COMPARISON_OPERATOR,
    NOT_OPERATOR,
    BOOLEAN_OPERATOR,
    AWAIT,
    LAMBDA,
    CONDITIONAL_EXPRESSION,
    NAMED_EXPRESSION,
    AS_PATTERN,
    EXPRESSION_LIST,
    LIST_SPLAT_OR_EXPRESSIONS,

    // `primary_expression`
    BINARY_OPERATOR,
    IDENTIFIER,
    KEYWORD_IDENTIFIER,
    STRING,
    CONCATENATED_STRING,
    INTEGER,
    FLOAT,
    TRUE,
    FALSE,
    NONE,
    UNARY_OPERATOR,
    ATTRIBUTE,
    SUBSCRIPT,
    CALL,
    LIST,
    LIST_COMPREHENSION,
    DICTIONARY,
    DICTIONARY_COMPREHENSION,
    SET,
    SET_COMPREHENSION,
    TUPLE,
    PARENTHESIZED_EXPRESSION,
    GENERATOR_EXPRESSION,
    ELLIPSIS,

    FOR_IN_CLAUSE,
    IF_CLAUSE,

    TYPED_PARAMETER,
    KEYWORD_ARGUMENT,
    DEFAULT_PARAMETER,
    TYPED_DEFAULT_PARAMETER,
    LIST_SPLAT_PATTERN,
    TUPLE_PATTERN,
    LIST_PATTERN,
    KEYWORD_SEPARATOR,
    POSITIONAL_SEPARATOR,
    DICTIONARY_SPLAT_PATTERN,
    PATTERN_LIST,

    COMMENT,
    MODULE,
    PARAMETERS,
    TYPE,
    BLOCK,

    // scaffolding
    PAIR,
    ARGUMENT_LIST,
    DECORATOR,

    // special function arguments
    LIST_SPLAT,
    DICTIONARY_SPLAT,
}

impl<'a> Production<'a> {
    pub fn new(kind: ProductionKind, node: &'a Node<'a>) -> Self {
        Production {
            production_kind: kind,
            node,
        }
    }
}

///
/// Wrap a sitter `Node` into its structured wrapper (`NodeType`)
///
pub fn get_node_type<'a>(node: &'a Node<'a>) -> NodeType<'a> {
    match node.kind() {
        // keywords
        "as" => NodeType::Keyword(Keyword::AS),
        "assert" => NodeType::Keyword(Keyword::ASSERT),
        "async" => NodeType::Keyword(Keyword::ASYNC),
        // "await" => NodeType::Keyword(Keyword::AWAIT),
        "break" => NodeType::Keyword(Keyword::BREAK),
        "class" => NodeType::Keyword(Keyword::CLASS),
        "continue" => NodeType::Keyword(Keyword::CONTINUE),
        "def" => NodeType::Keyword(Keyword::DEF),
        "del" => NodeType::Keyword(Keyword::DEL),
        "elif" => NodeType::Keyword(Keyword::ELIF),
        "else" => NodeType::Keyword(Keyword::ELSE),
        "except" => NodeType::Keyword(Keyword::EXCEPT),
        "finally" => NodeType::Keyword(Keyword::FINALLY),
        "for" => NodeType::Keyword(Keyword::FOR),
        "from" => NodeType::Keyword(Keyword::FROM),
        "global" => NodeType::Keyword(Keyword::GLOBAL),
        "if" => NodeType::Keyword(Keyword::IF),
        "import" => NodeType::Keyword(Keyword::IMPORT),
        "in" => NodeType::Keyword(Keyword::IN),
        "is" => NodeType::Keyword(Keyword::IS),
        "is not" => NodeType::Keyword(Keyword::IS_NOT),
        "not in" => NodeType::Keyword(Keyword::NOT_IN),
        // "lambda" => NodeType::Keyword(Keyword::LAMBDA),
        "None" => NodeType::Keyword(Keyword::NONE),
        "nonlocal" => NodeType::Keyword(Keyword::NONLOCAL),
        "not" => NodeType::Keyword(Keyword::NOT),
        "pass" => NodeType::Keyword(Keyword::PASS),
        "raise" => NodeType::Keyword(Keyword::RAISE),
        "return" => NodeType::Keyword(Keyword::RETURN),
        "try" => NodeType::Keyword(Keyword::TRY),
        "while" => NodeType::Keyword(Keyword::WHILE),
        "with" => NodeType::Keyword(Keyword::WITH),

        // boolean operator
        "and" => NodeType::Keyword(Keyword::AND),
        "or" => NodeType::Keyword(Keyword::OR),

        // operator
        "==" => NodeType::ComparisonOperator(ComparisonOperator::EQUAL),
        ">" => NodeType::ComparisonOperator(ComparisonOperator::GREATER_THAN),
        ">=" => NodeType::ComparisonOperator(ComparisonOperator::GREATER_THAN_EQUAL),
        "<" => NodeType::ComparisonOperator(ComparisonOperator::LESS_THAN),
        "<=" => NodeType::ComparisonOperator(ComparisonOperator::LESS_THAN_EQUAL),
        "!=" => NodeType::ComparisonOperator(ComparisonOperator::NOT_EQUAL),

        ":=" => NodeType::BinaryOperator(BinaryOperator::ASSIGN_EXPR),
        "@" => NodeType::BinaryOperator(BinaryOperator::AT),
        "&" => NodeType::BinaryOperator(BinaryOperator::BITWISE_AND),
        "|" => NodeType::BinaryOperator(BinaryOperator::BITWISE_OR),
        "^" => NodeType::BinaryOperator(BinaryOperator::BITWISE_XOR),
        "//" => NodeType::BinaryOperator(BinaryOperator::DOUBLE_SLASH),
        "**" => NodeType::BinaryOperator(BinaryOperator::DOUBLE_STAR),
        "<<" => NodeType::BinaryOperator(BinaryOperator::LEFT_SHIFT),
        "-" => NodeType::BinaryOperator(BinaryOperator::MINUS),
        "%" => NodeType::BinaryOperator(BinaryOperator::PERCENT),
        "+" => NodeType::BinaryOperator(BinaryOperator::PLUS),
        ">>" => NodeType::BinaryOperator(BinaryOperator::RIGHT_SHIFT),
        "/" => NodeType::BinaryOperator(BinaryOperator::SLASH),
        "*" => NodeType::BinaryOperator(BinaryOperator::STAR),
        "~" => NodeType::BinaryOperator(BinaryOperator::TILDE),

        //aug assignment operators
        "@=" => NodeType::AugAssignOperator(AugAssignOperator::AT_EQUAL),
        "&=" => NodeType::AugAssignOperator(AugAssignOperator::BITWISE_AND_EQUAL),
        "|=" => NodeType::AugAssignOperator(AugAssignOperator::BITWISE_OR_EQUAL),
        "//=" => NodeType::AugAssignOperator(AugAssignOperator::DOUBLE_SLASH_EQUAL),
        "**=" => NodeType::AugAssignOperator(AugAssignOperator::DOUBLE_STAR_EQUAL),
        "<<=" => NodeType::AugAssignOperator(AugAssignOperator::LEFT_SHIFT_EQUAL),
        "-=" => NodeType::AugAssignOperator(AugAssignOperator::MINUS_EQUAL),
        "%=" => NodeType::AugAssignOperator(AugAssignOperator::PERCENT_EQUAL),
        "+=" => NodeType::AugAssignOperator(AugAssignOperator::PLUS_EQUAL),
        ">>=" => NodeType::AugAssignOperator(AugAssignOperator::RIGHT_SHIFT_EQUAL),
        "/=" => NodeType::AugAssignOperator(AugAssignOperator::SLASH_EQUAL),
        "*=" => NodeType::AugAssignOperator(AugAssignOperator::STAR_EQUAL),
        "^=" => NodeType::AugAssignOperator(AugAssignOperator::BITWISE_XOR_EQUAL),

        // delimiter
        "->" => NodeType::Delimiter(Delimiter::ARROW),
        "=" => NodeType::Delimiter(Delimiter::ASSIGN),
        "\\" => NodeType::Delimiter(Delimiter::BACK_SLASH),
        ":" => NodeType::Delimiter(Delimiter::COLON),
        "," => NodeType::Delimiter(Delimiter::COMMA),
        "\"" => NodeType::Delimiter(Delimiter::DUBLE_QUOTE),
        "#" => NodeType::Delimiter(Delimiter::HASH),
        "[" => NodeType::Delimiter(Delimiter::LEFT_BRACKET),
        "{" => NodeType::Delimiter(Delimiter::LEFT_CURLY),
        "(" => NodeType::Delimiter(Delimiter::LEFT_PAREN),
        "." => NodeType::Delimiter(Delimiter::PERIOD),
        "]" => NodeType::Delimiter(Delimiter::RIGHT_BRACKET),
        "}" => NodeType::Delimiter(Delimiter::RIGHT_CURLY),
        ")" => NodeType::Delimiter(Delimiter::RIGHT_PAREN),
        ";" => NodeType::Delimiter(Delimiter::SEMICOLON),
        "'" => NodeType::Delimiter(Delimiter::SINGLE_QUOTE),

        // grammar productions
        "comment" => NodeType::Production(Production::new(ProductionKind::COMMENT, node)),
        "block" => NodeType::Production(Production::new(ProductionKind::BLOCK, node)),

        // `_simple_statement`
        "future_import_statement" => NodeType::Production(Production::new(
            ProductionKind::FUTURE_IMPORT_STATEMENT,
            node,
        )),
        "import_statement" => {
            NodeType::Production(Production::new(ProductionKind::IMPORT_STATEMENT, node))
        }
        "import_from_statement" => {
            NodeType::Production(Production::new(ProductionKind::IMPORT_FROM_STATEMENT, node))
        }
        "print_statement" => {
            NodeType::Production(Production::new(ProductionKind::PRINT_STATEMENT, node))
        }
        "assert_statement" => {
            NodeType::Production(Production::new(ProductionKind::ASSERT_STATEMENT, node))
        }
        "expression_statement" => {
            NodeType::Production(Production::new(ProductionKind::EXPRESSION_STATEMENT, node))
        }
        "return_statement" => {
            NodeType::Production(Production::new(ProductionKind::RETURN_STATEMENT, node))
        }
        "delete_statement" => {
            NodeType::Production(Production::new(ProductionKind::DELETE_STATEMENT, node))
        }
        "raise_statement" => {
            NodeType::Production(Production::new(ProductionKind::RAISE_STATEMENT, node))
        }
        "pass_statement" => {
            NodeType::Production(Production::new(ProductionKind::PASS_STATEMENT, node))
        }
        "break_statement" => {
            NodeType::Production(Production::new(ProductionKind::BREAK_STATEMENT, node))
        }
        "continue_statement" => {
            NodeType::Production(Production::new(ProductionKind::CONTINUE_STATEMENT, node))
        }
        "global_statement" => {
            NodeType::Production(Production::new(ProductionKind::GLOBAL_STATEMENT, node))
        }
        "nonlocal_statement" => {
            NodeType::Production(Production::new(ProductionKind::NONLOCAL_STATEMENT, node))
        }
        "exec_statement" => {
            NodeType::Production(Production::new(ProductionKind::EXEC_STATEMENT, node))
        }
        "if_statement" => NodeType::Production(Production::new(ProductionKind::IF_STATEMENT, node)),
        "for_statement" => {
            NodeType::Production(Production::new(ProductionKind::FOR_STATEMENT, node))
        }
        "while_statement" => {
            NodeType::Production(Production::new(ProductionKind::WHILE_STATEMENT, node))
        }
        "try_statement" => {
            NodeType::Production(Production::new(ProductionKind::TRY_STATEMENT, node))
        }
        "with_statement" => {
            NodeType::Production(Production::new(ProductionKind::WITH_STATEMENT, node))
        }
        "function_definition" => {
            NodeType::Production(Production::new(ProductionKind::FUNCTION_DEFINITION, node))
        }
        "class_definition" => {
            NodeType::Production(Production::new(ProductionKind::CLASS_DEFINITION, node))
        }
        "decorated_definition" => {
            NodeType::Production(Production::new(ProductionKind::DECORATED_DEFINITION, node))
        }
        "match_statement" => {
            NodeType::Production(Production::new(ProductionKind::MATCH_STATEMENT, node))
        }

        // `expression_statement`
        "assignment" => NodeType::Production(Production::new(ProductionKind::ASSIGNMENT, node)),
        "augmented_assignment" => {
            NodeType::Production(Production::new(ProductionKind::AUGMENTED_ASSIGNMENT, node))
        }
        "yield" => NodeType::Production(Production::new(ProductionKind::YIELD, node)),

        // `expression`
        "comparison_operator" => {
            NodeType::Production(Production::new(ProductionKind::COMPARISON_OPERATOR, node))
        }
        "not_operator" => NodeType::Production(Production::new(ProductionKind::NOT_OPERATOR, node)),
        "boolean_operator" => {
            NodeType::Production(Production::new(ProductionKind::BOOLEAN_OPERATOR, node))
        }
        "await" => NodeType::Production(Production::new(ProductionKind::AWAIT, node)),
        "lambda" => NodeType::Production(Production::new(ProductionKind::LAMBDA, node)),
        "conditional_expression" => NodeType::Production(Production::new(
            ProductionKind::CONDITIONAL_EXPRESSION,
            node,
        )),
        "named_expression" => {
            NodeType::Production(Production::new(ProductionKind::NAMED_EXPRESSION, node))
        }
        "as_pattern" => NodeType::Production(Production::new(ProductionKind::AS_PATTERN, node)),

        // `primary_expression`
        "binary_operator" => {
            NodeType::Production(Production::new(ProductionKind::BINARY_OPERATOR, node))
        }
        "identifier" => NodeType::Production(Production::new(ProductionKind::IDENTIFIER, node)),
        "keyword_identifier" => {
            NodeType::Production(Production::new(ProductionKind::KEYWORD_IDENTIFIER, node))
        }
        "string" => NodeType::Production(Production::new(ProductionKind::STRING, node)),
        "concatenated_string" => {
            NodeType::Production(Production::new(ProductionKind::CONCATENATED_STRING, node))
        }
        "integer" => NodeType::Production(Production::new(ProductionKind::INTEGER, node)),
        "float" => NodeType::Production(Production::new(ProductionKind::FLOAT, node)),
        "true" => NodeType::Production(Production::new(ProductionKind::TRUE, node)),
        "false" => NodeType::Production(Production::new(ProductionKind::FALSE, node)),
        "none" => NodeType::Production(Production::new(ProductionKind::NONE, node)),
        "unary_operator" => {
            NodeType::Production(Production::new(ProductionKind::UNARY_OPERATOR, node))
        }
        "attribute" => NodeType::Production(Production::new(ProductionKind::ATTRIBUTE, node)),
        "subscript" => NodeType::Production(Production::new(ProductionKind::SUBSCRIPT, node)),
        "call" => NodeType::Production(Production::new(ProductionKind::CALL, node)),
        "list" => NodeType::Production(Production::new(ProductionKind::LIST, node)),
        "list_comprehension" => {
            NodeType::Production(Production::new(ProductionKind::LIST_COMPREHENSION, node))
        }
        "dictionary" => NodeType::Production(Production::new(ProductionKind::DICTIONARY, node)),
        "dictionary_comprehension" => NodeType::Production(Production::new(
            ProductionKind::DICTIONARY_COMPREHENSION,
            node,
        )),
        "set" => NodeType::Production(Production::new(ProductionKind::SET, node)),
        "set_comprehension" => {
            NodeType::Production(Production::new(ProductionKind::SET_COMPREHENSION, node))
        }

        "for_in_clause" => {
            NodeType::Production(Production::new(ProductionKind::FOR_IN_CLAUSE, node))
        }
        "if_clause" => NodeType::Production(Production::new(ProductionKind::IF_CLAUSE, node)),

        "tuple" => NodeType::Production(Production::new(ProductionKind::TUPLE, node)),
        "parenthesized_expression" => NodeType::Production(Production::new(
            ProductionKind::PARENTHESIZED_EXPRESSION,
            node,
        )),
        "generator_expression" => {
            NodeType::Production(Production::new(ProductionKind::GENERATOR_EXPRESSION, node))
        }
        "ellipsis" => NodeType::Production(Production::new(ProductionKind::ELLIPSIS, node)),

        "argument_list" => {
            NodeType::Production(Production::new(ProductionKind::ARGUMENT_LIST, node))
        }
        "parameters" => NodeType::Production(Production::new(ProductionKind::PARAMETERS, node)),

        "pair" => NodeType::Production(Production::new(ProductionKind::PAIR, node)),

        // special function arguments
        "list_splat" => NodeType::Production(Production::new(ProductionKind::LIST_SPLAT, node)),
        "dictionary_splat" => {
            NodeType::Production(Production::new(ProductionKind::DICTIONARY_SPLAT, node))
        }

        // `parameter`
        "typed_parameter" => {
            NodeType::Production(Production::new(ProductionKind::TYPED_PARAMETER, node))
        }
        "keyword_argument" => {
            NodeType::Production(Production::new(ProductionKind::KEYWORD_ARGUMENT, node))
        }
        "default_parameter" => {
            NodeType::Production(Production::new(ProductionKind::DEFAULT_PARAMETER, node))
        }
        "typed_default_parameter" => NodeType::Production(Production::new(
            ProductionKind::TYPED_DEFAULT_PARAMETER,
            node,
        )),
        "list_splat_pattern" => {
            NodeType::Production(Production::new(ProductionKind::LIST_SPLAT_PATTERN, node))
        }
        "tuple_pattern" => {
            NodeType::Production(Production::new(ProductionKind::TUPLE_PATTERN, node))
        }
        "list_pattern" => NodeType::Production(Production::new(ProductionKind::LIST_PATTERN, node)),
        "pattern_list" => NodeType::Production(Production::new(ProductionKind::PATTERN_LIST, node)),
        "keyword_separator" => {
            NodeType::Production(Production::new(ProductionKind::KEYWORD_SEPARATOR, node))
        }
        "positional_separator" => {
            NodeType::Production(Production::new(ProductionKind::POSITIONAL_SEPARATOR, node))
        }
        "expression_list" => {
            NodeType::Production(Production::new(ProductionKind::EXPRESSION_LIST, node))
        }
        "list_splat_or_expressions" => NodeType::Production(Production::new(
            ProductionKind::LIST_SPLAT_OR_EXPRESSIONS,
            node,
        )),
        "dictionary_splat_pattern" => NodeType::Production(Production::new(
            ProductionKind::DICTIONARY_SPLAT_PATTERN,
            node,
        )),

        // place production below properly
        "module" => NodeType::Production(Production::new(ProductionKind::MODULE, node)),
        "type" => NodeType::Production(Production::new(ProductionKind::TYPE, node)),

        "decorator" => NodeType::Production(Production::new(ProductionKind::DECORATOR, node)),

        "ERROR" => NodeType::Error,

        _ => panic!(
            "unmatched keyword node type: {}[{}]",
            node.kind(),
            node.is_named()
        ),
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree

#[derive(Debug)]
#[derive(PartialEq)]
#[derive(Default)]
pub struct StringCategory {
    pub is_byte: bool,
    pub is_raw: bool,
    pub is_format: bool,
    pub is_unicode: bool,
}

#[derive(Debug)]
#[derive(PartialEq)]
pub struct InvalidStringCategoryError {
    pub invalid_prefix: String,
}

/// Return a string containing all characters prior to the first ' or "
pub fn string_prefix(string_contents: &str) -> String {
    let mut prefix: String = String::from("");
    for ch in string_contents.chars() {
        if ch == '\'' || ch == '"' {
            break;
        }
        prefix.push(ch);
    }
    prefix
}

/// categorize_string will return a StringCategory
/// containing the type of string which has been passed to the function
/// i.e. a byte, format, unicode or raw string as signified by the prefix
/// of the string, e.g. `b"a byte string"` is a byte string etc.
/// If an invalid combination of prefix is provided then this is returned
/// as a InvalidStringCategoryError for processing by the caller
pub fn categorize_string(
    to_categorize: &str,
) -> Result<StringCategory, InvalidStringCategoryError> {
    // Only valid prefixes are:
    // "r" | "u" | "f" | "b"
    // "fr" | "rf"
    // "br"| "rb"
    let mut prefix = string_prefix(to_categorize);

    prefix = prefix.to_lowercase();
    match prefix.len() {
        1 => match prefix.as_str() {
            "r" => Ok(StringCategory {
                is_raw: true,
                ..Default::default()
            }),
            "u" => Ok(StringCategory {
                is_unicode: true,
                ..Default::default()
            }),
            "f" => Ok(StringCategory {
                is_format: true,
                ..Default::default()
            }),
            "b" => Ok(StringCategory {
                is_byte: true,
                ..Default::default()
            }),
            _ => Err(InvalidStringCategoryError {
                invalid_prefix: prefix,
            }),
        },
        2 => match prefix.as_str() {
            "fr" | "rf" => Ok(StringCategory {
                is_format: true,
                is_raw: true,
                ..Default::default()
            }),
            "br" | "rb" => Ok(StringCategory {
                is_byte: true,
                is_raw: true,
                ..Default::default()
            }),
            _ => Err(InvalidStringCategoryError {
                invalid_prefix: prefix,
            }),
        },
        0 => Ok(StringCategory {
            ..Default::default()
        }),
        _ => Err(InvalidStringCategoryError {
            invalid_prefix: prefix,
        }),
    }
}