<commit_before>// This parses the ops.pbtxt file in TensorFlow and returns a list of op definitions. This is
// currently implemented directly with nom to parse the text proto, but ideally this would use a
// proto library with support for parsing text protos.

use crate::protos::AttrValue;
use crate::protos::AttrValue_ListValue;
use crate::protos::DataType;
use crate::protos::FullTypeDef;
use crate::protos::FullTypeId;
use crate::protos::OpDef;
use crate::protos::OpDef_ArgDef;
use crate::protos::OpDef_AttrDef;
use crate::protos::OpDeprecation;
use crate::protos::TensorProto;
use crate::protos::TensorShapeProto;
use crate::protos::TensorShapeProto_Dim;
use nom::branch::alt;
use nom::bytes::complete::tag;
use nom::character::complete::anychar;
use nom::character::complete::char;
use nom::character::complete::none_of;
use nom::character::complete::one_of;
use nom::combinator::cut;
use nom::combinator::map;
use nom::combinator::map_res;
use nom::error::context;
use nom::error::make_error;
use nom::error::ErrorKind;
use nom::error::VerboseError;
use nom::multi::many0;
use nom::multi::many1;
use nom::multi::many_till;
use nom::multi::separated_list0;
use nom::sequence::delimited;
use nom::sequence::pair;
use nom::sequence::preceded;
use nom::sequence::terminated;
use nom::IResult;
use protobuf::Message;
use protobuf::ProtobufError;
use std::error::Error;
use std::fmt;
use std::fmt::Display;
use std::result::Result;

fn merge_protos<M: Message>(ops: Vec<M>) -> Result<M, ProtobufError> {
    let mut out = Vec::new();
    for op in ops {
        op.write_to_vec(&mut out)?;
    }
    Message::parse_from_bytes(&out)
}

type ParseResult<'a, O> = IResult<&'a [u8], O, VerboseError<&'a [u8]>>;

fn space<'a>() -> impl FnMut(&'a [u8]) -> ParseResult<'a, ()> {
    map(many0(one_of(" \t\r\n")), |_| ())
}

fn eof<'a>() -> impl Fn(&'a [u8]) -> ParseResult<'a, ()> {
    move |input: &[u8]| {
        if input.len() == 0 {
            Ok((input, ()))
        } else {
            Err(nom::Err::Error(make_error(input, ErrorKind::Eof)))
        }
    }
}

fn string<'a>(input: &'a [u8]) -> ParseResult<'a, String> {
    let string_char = alt((
        none_of("\"\\"),
        map(tag("\\n"), |_| '\n'),
        map(tag("\\t"), |_| '\t'),
        map(tag("\\\""), |_| '"'),
        map(tag("\\'"), |_| '\''),
    ));
    delimited(
        char('"'),
        map(many0(string_char), |v| v.into_iter().collect()),
        char('"'),
    )(input)
}

fn boolean<'a>(input: &'a [u8]) -> ParseResult<'a, bool> {
    alt((map(tag("true"), |_| true), map(tag("false"), |_| false)))(input)
}

fn int64<'a>(input: &'a [u8]) -> ParseResult<'a, i64> {
    map_res(many1(one_of("0123456789+-")), |v| {
        str::parse::<i64>(&v.iter().collect::<String>())
    })(input)
}

fn int32<'a>(input: &'a [u8]) -> ParseResult<'a, i32> {
    map_res(many1(one_of("0123456789+-")), |v| {
        str::parse::<i32>(&v.iter().collect::<String>())
    })(input)
}

fn float_<'a>(input: &'a [u8]) -> ParseResult<'a, f32> {
    map_res(many1(one_of("0123456789+-.aefinAEFIN")), |v| {
        str::parse::<f32>(&v.iter().collect::<String>())
    })(input)
}

fn identifier<'a>(input: &'a [u8]) -> ParseResult<'a, String> {
    let identifier_start = map_res(anychar, |c| match c {
        'A'..='Z' | 'a'..='z' | '_' => Ok(c),
        _ => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
            input,
            ErrorKind::OneOf,
        ))),
    });
    let identifier_part = map_res(anychar, |c| match c {
        'A'..='Z' | 'a'..='z' | '_' | '0'..='9' => Ok(c),
        _ => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
            input,
            ErrorKind::OneOf,
        ))),
    });
    map(
        pair(identifier_start, many0(identifier_part)),
        |(first, rest)| [first].iter().chain(rest.iter()).collect::<String>(),
    )(input)
}

fn data_type<'a>(input: &'a [u8]) -> ParseResult<'a, DataType> {
    map_res(identifier, |s| match &s as &str {
        "DT_BFLOAT16" => Ok(DataType::DT_BFLOAT16),
        "DT_HALF" => Ok(DataType::DT_HALF),
        "DT_FLOAT" => Ok(DataType::DT_FLOAT),
        "DT_DOUBLE" => Ok(DataType::DT_DOUBLE),
        "DT_COMPLEX64" => Ok(DataType::DT_COMPLEX64),
        "DT_COMPLEX128" => Ok(DataType::DT_COMPLEX128),
        "DT_INT8" => Ok(DataType::DT_INT8),
        "DT_INT16" => Ok(DataType::DT_INT16),
        "DT_INT32" => Ok(DataType::DT_INT32),
        "DT_INT64" => Ok(DataType::DT_INT64),
        "DT_UINT8" => Ok(DataType::DT_UINT8),
        "DT_UINT16" => Ok(DataType::DT_UINT16),
        "DT_UINT32" => Ok(DataType::DT_UINT32),
        "DT_UINT64" => Ok(DataType::DT_UINT64),
        "DT_QINT8" => Ok(DataType::DT_QINT8),
        "DT_QINT16" => Ok(DataType::DT_QINT16),
        "DT_QINT32" => Ok(DataType::DT_QINT32),
        "DT_QUINT8" => Ok(DataType::DT_QUINT8),
        "DT_QUINT16" => Ok(DataType::DT_QUINT16),
        "DT_STRING" => Ok(DataType::DT_STRING),
        "DT_VARIANT" => Ok(DataType::DT_VARIANT),
        "DT_BOOL" => Ok(DataType::DT_BOOL),
        "DT_RESOURCE" => Ok(DataType::DT_RESOURCE),
        _ => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
            input,
            ErrorKind::Alt,
        ))),
    })(input)
}

fn type_id<'a>(input: &'a [u8]) -> ParseResult<'a, FullTypeId> {
    map_res(identifier, |s| match &s as &str {
        "TFT_UNSET" => Ok(FullTypeId::TFT_UNSET),
        "TFT_VAR" => Ok(FullTypeId::TFT_VAR),
        "TFT_ANY" => Ok(FullTypeId::TFT_ANY),
        "TFT_PRODUCT" => Ok(FullTypeId::TFT_PRODUCT),
        "TFT_CALLABLE" => Ok(FullTypeId::TFT_CALLABLE),
        "TFT_TENSOR" => Ok(FullTypeId::TFT_TENSOR),
        "TFT_ARRAY" => Ok(FullTypeId::TFT_ARRAY),
        "TFT_OPTIONAL" => Ok(FullTypeId::TFT_OPTIONAL),
        "TFT_DATASET2" => Ok(FullTypeId::TFT_OPTIONAL),
        "TFT_BOOL" => Ok(FullTypeId::TFT_OPTIONAL),
        "TFT_UINT8" => Ok(FullTypeId::TFT_OPTIONAL),
        "TFT_UINT16" => Ok(FullTypeId::TFT_OPTIONAL),
        "TFT_UINT32" => Ok(FullTypeId::TFT_OPTIONAL),
        "TFT_UINT64" => Ok(FullTypeId::TFT_OPTIONAL),
        "TFT_INT8" => Ok(FullTypeId::TFT_INT8),
        "TFT_INT16" => Ok(FullTypeId::TFT_INT16),
        "TFT_INT32" => Ok(FullTypeId::TFT_INT32),
        "TFT_INT64" => Ok(FullTypeId::TFT_INT64),
        "TFT_HALF" => Ok(FullTypeId::TFT_HALF),
        "TFT_FLOAT" => Ok(FullTypeId::TFT_FLOAT),
        "TFT_DOUBLE" => Ok(FullTypeId::TFT_DOUBLE),
        "TFT_BFLOAT16" => Ok(FullTypeId::TFT_BFLOAT16),
        "TFT_COMPLEX64" => Ok(FullTypeId::TFT_COMPLEX64),
        "TFT_COMPLEX128" => Ok(FullTypeId::TFT_COMPLEX128),
        "TFT_STRING" => Ok(FullTypeId::TFT_STRING),
        _ => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
            input,
            ErrorKind::Alt,
        ))),
    })(input)
}

fn scalar_field<'a, M, T, Parser, ParserFactory, S>(
    name: &'a str,
    parser_factory: ParserFactory,
    setter: S,
) -> impl Fn(&'a [u8]) -> ParseResult<'a, M>
where
    M: Message,
    Parser: Fn(&'a [u8]) -> ParseResult<'a, T>,
    ParserFactory: Fn() -> Parser,<commit_msg>[codegen] simplify `type_id` extraction.<commit_after>// This parses the ops.pbtxt file in TensorFlow and returns a list of op definitions. This is
// currently implemented directly with nom to parse the text proto, but ideally this would use a
// proto library with support for parsing text protos.

use crate::protos::AttrValue;
use crate::protos::AttrValue_ListValue;
use crate::protos::DataType;
use crate::protos::FullTypeDef;
use crate::protos::FullTypeId;
use crate::protos::OpDef;
use crate::protos::OpDef_ArgDef;
use crate::protos::OpDef_AttrDef;
use crate::protos::OpDeprecation;
use crate::protos::TensorProto;
use crate::protos::TensorShapeProto;
use crate::protos::TensorShapeProto_Dim;
use nom::branch::alt;
use nom::bytes::complete::tag;
use nom::character::complete::anychar;
use nom::character::complete::char;
use nom::character::complete::none_of;
use nom::character::complete::one_of;
use nom::combinator::cut;
use nom::combinator::map;
use nom::combinator::map_res;
use nom::error::context;
use nom::error::make_error;
use nom::error::ErrorKind;
use nom::error::VerboseError;
use nom::multi::many0;
use nom::multi::many1;
use nom::multi::many_till;
use nom::multi::separated_list0;
use nom::sequence::delimited;
use nom::sequence::pair;
use nom::sequence::preceded;
use nom::sequence::terminated;
use nom::IResult;
use protobuf::Message;
use protobuf::ProtobufEnum;
use protobuf::ProtobufError;
use std::error::Error;
use std::fmt;
use std::fmt::Display;
use std::result::Result;

fn merge_protos<M: Message>(ops: Vec<M>) -> Result<M, ProtobufError> {
    let mut out = Vec::new();
    for op in ops {
        op.write_to_vec(&mut out)?;
    }
    Message::parse_from_bytes(&out)
}

type ParseResult<'a, O> = IResult<&'a [u8], O, VerboseError<&'a [u8]>>;

fn space<'a>() -> impl FnMut(&'a [u8]) -> ParseResult<'a, ()> {
    map(many0(one_of(" \t\r\n")), |_| ())
}

fn eof<'a>() -> impl Fn(&'a [u8]) -> ParseResult<'a, ()> {
    move |input: &[u8]| {
        if input.len() == 0 {
            Ok((input, ()))
        } else {
            Err(nom::Err::Error(make_error(input, ErrorKind::Eof)))
        }
    }
}

fn string<'a>(input: &'a [u8]) -> ParseResult<'a, String> {
    let string_char = alt((
        none_of("\"\\"),
        map(tag("\\n"), |_| '\n'),
        map(tag("\\t"), |_| '\t'),
        map(tag("\\\""), |_| '"'),
        map(tag("\\'"), |_| '\''),
    ));
    delimited(
        char('"'),
        map(many0(string_char), |v| v.into_iter().collect()),
        char('"'),
    )(input)
}

fn boolean<'a>(input: &'a [u8]) -> ParseResult<'a, bool> {
    alt((map(tag("true"), |_| true), map(tag("false"), |_| false)))(input)
}

fn int64<'a>(input: &'a [u8]) -> ParseResult<'a, i64> {
    map_res(many1(one_of("0123456789+-")), |v| {
        str::parse::<i64>(&v.iter().collect::<String>())
    })(input)
}

fn int32<'a>(input: &'a [u8]) -> ParseResult<'a, i32> {
    map_res(many1(one_of("0123456789+-")), |v| {
        str::parse::<i32>(&v.iter().collect::<String>())
    })(input)
}

fn float_<'a>(input: &'a [u8]) -> ParseResult<'a, f32> {
    map_res(many1(one_of("0123456789+-.aefinAEFIN")), |v| {
        str::parse::<f32>(&v.iter().collect::<String>())
    })(input)
}

fn identifier<'a>(input: &'a [u8]) -> ParseResult<'a, String> {
    let identifier_start = map_res(anychar, |c| match c {
        'A'..='Z' | 'a'..='z' | '_' => Ok(c),
        _ => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
            input,
            ErrorKind::OneOf,
        ))),
    });
    let identifier_part = map_res(anychar, |c| match c {
        'A'..='Z' | 'a'..='z' | '_' | '0'..='9' => Ok(c),
        _ => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
            input,
            ErrorKind::OneOf,
        ))),
    });
    map(
        pair(identifier_start, many0(identifier_part)),
        |(first, rest)| [first].iter().chain(rest.iter()).collect::<String>(),
    )(input)
}

fn data_type<'a>(input: &'a [u8]) -> ParseResult<'a, DataType> {
    map_res(identifier, |s| match &s as &str {
        "DT_BFLOAT16" => Ok(DataType::DT_BFLOAT16),
        "DT_HALF" => Ok(DataType::DT_HALF),
        "DT_FLOAT" => Ok(DataType::DT_FLOAT),
        "DT_DOUBLE" => Ok(DataType::DT_DOUBLE),
        "DT_COMPLEX64" => Ok(DataType::DT_COMPLEX64),
        "DT_COMPLEX128" => Ok(DataType::DT_COMPLEX128),
        "DT_INT8" => Ok(DataType::DT_INT8),
        "DT_INT16" => Ok(DataType::DT_INT16),
        "DT_INT32" => Ok(DataType::DT_INT32),
        "DT_INT64" => Ok(DataType::DT_INT64),
        "DT_UINT8" => Ok(DataType::DT_UINT8),
        "DT_UINT16" => Ok(DataType::DT_UINT16),
        "DT_UINT32" => Ok(DataType::DT_UINT32),
        "DT_UINT64" => Ok(DataType::DT_UINT64),
        "DT_QINT8" => Ok(DataType::DT_QINT8),
        "DT_QINT16" => Ok(DataType::DT_QINT16),
        "DT_QINT32" => Ok(DataType::DT_QINT32),
        "DT_QUINT8" => Ok(DataType::DT_QUINT8),
        "DT_QUINT16" => Ok(DataType::DT_QUINT16),
        "DT_STRING" => Ok(DataType::DT_STRING),
        "DT_VARIANT" => Ok(DataType::DT_VARIANT),
        "DT_BOOL" => Ok(DataType::DT_BOOL),
        "DT_RESOURCE" => Ok(DataType::DT_RESOURCE),
        _ => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
            input,
            ErrorKind::Alt,
        ))),
    })(input)
}

fn type_id<'a>(input: &'a [u8]) -> ParseResult<'a, FullTypeId> {
    map_res(identifier, |s| {
        match FullTypeId::values()
            .iter()
            .filter(|d| d.descriptor().name() == s)
            .next()
        {
            Some(d) => Ok(FullTypeId::from_i32(d.value()).unwrap()),
            None => Err(nom::Err::<VerboseError<&'a [u8]>>::Error(make_error(
                input,
                ErrorKind::Alt,
            ))),
        }
    })(input)
}

fn scalar_field<'a, M, T, Parser, ParserFactory, S>(
    name: &'a str,
    parser_factory: ParserFactory,
    setter: S,
) -> impl Fn(&'a [u8]) -> ParseResult<'a, M>
where
    M: Message,
    Parser: Fn(&'a [u8]) -> ParseResult<'a, T>,
    ParserFactory: Fn() -> Parser,
{'fraction_non_alphanumeric': 0.21140377588306944, 'fraction_numerical': 0.023447015834348356, 'mean_word_length': 3.5614583333333334, 'pattern_counts': {'":': 0, '<': 91, '<?xml version=': 0, '>': 209, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '652913', 'n_tokens_mistral': 5138, 'n_tokens_neox': 5091, 'n_words': 826}