use std::fmt::Display;
use proc_macro2::{Delimiter, TokenTree};
use syn::parse::{Parse, ParseStream};
use syn::{Error, Ident, Result, Token, braced, bracketed, parenthesized};
use crate::model::header::{RicciAliasDeclaration, RicciIndexDeclaration, RicciIndexer};
use phf::{Set, phf_set};
static ILLEGAL_KEYWORDS: Set<&'static str> = phf_set! {
"while",
"loop",
"break",
"continue",
"return",
"fn",
"struct",
"enum",
"type",
"trait",
"mod",
"use",
"extern",
};
pub enum RicciSegment {
TensorCall {
tensor_name: Ident,
indexers: Vec<RicciIndexer>,
},
SubLambda(Box<RicciLambda>),
SubGroup {
delimiter: Delimiter,
group: Box<RicciGroup>,
},
Token(TokenTree),
}
pub struct RicciGroup {
pub segments: Vec<RicciSegment>,
}
pub struct RicciFilter {
pub if_keyword: Token![if],
pub segments: Vec<RicciSegment>,
}
pub struct RicciLambda {
pub index_declaration: RicciIndexDeclaration,
pub alias_declarations: Vec<RicciAliasDeclaration>,
pub filter: Option<RicciFilter>,
pub body: RicciGroup,
}
impl Parse for RicciSegment {
fn parse(input: ParseStream) -> Result<Self> {
if input.peek(syn::token::Paren) {
let content;
parenthesized!(content in input);
let group: Box<RicciGroup> = content.parse()?;
Ok(RicciSegment::SubGroup {
group,
delimiter: Delimiter::Parenthesis,
})
} else if input.peek(syn::token::Brace) {
let content;
braced!(content in input);
let group: Box<RicciGroup> = content.parse()?;
Ok(RicciSegment::SubGroup {
group,
delimiter: Delimiter::Brace,
})
} else if input.peek(syn::token::Bracket) {
let content;
bracketed!(content in input);
let group: Box<RicciGroup> = content.parse()?;
Ok(RicciSegment::SubGroup {
group,
delimiter: Delimiter::Bracket,
})
} else if input.peek(Token![for]) {
let lambda: Box<RicciLambda> = input.parse()?;
Ok(RicciSegment::SubLambda(lambda))
} else if input.peek(Ident) && input.peek2(syn::token::Bracket) {
let tensor_name: Ident = input.parse()?;
let content;
bracketed!(content in input);
let mut indexers: Vec<RicciIndexer> = Vec::new();
let mut first = true;
while !content.is_empty() {
if first {
first = false;
} else {
content.parse::<Token![,]>()?;
}
let indexer: RicciIndexer = content.parse()?;
indexers.push(indexer);
}
Ok(RicciSegment::TensorCall {
tensor_name,
indexers,
})
} else {
let token: TokenTree = input.parse()?;
if let TokenTree::Ident(ident) = &token
&& ILLEGAL_KEYWORDS.contains(ident.to_string().as_str())
{
return Err(Error::new(
ident.span(),
format!("Keyword {} is illegal.", ident),
));
}
Ok(RicciSegment::Token(token))
}
}
}
impl Parse for RicciGroup {
fn parse(input: ParseStream) -> Result<Self> {
let mut segments = Vec::<RicciSegment>::new();
while !input.is_empty() {
let segment = input.parse::<RicciSegment>()?;
segments.push(segment);
}
Ok(Self { segments })
}
}
impl Parse for RicciFilter {
fn parse(input: ParseStream) -> Result<Self> {
let if_keyword = input.parse::<Token![if]>()?;
let mut segments = Vec::<RicciSegment>::new();
while !input.is_empty() && !input.peek(Token![=>]) {
let segment = input.parse::<RicciSegment>()?;
segments.push(segment);
}
Ok(Self {
if_keyword,
segments,
})
}
}
impl Parse for RicciLambda {
fn parse(input: ParseStream) -> Result<Self> {
let index_declaration: RicciIndexDeclaration = input.parse()?;
let mut alias_declarations: Vec<RicciAliasDeclaration> = Vec::new();
while input.peek(Token![let]) {
let alias_declaration: RicciAliasDeclaration = input.parse()?;
alias_declarations.push(alias_declaration);
}
let filter = if input.peek(Token![if]) {
Some(input.parse::<RicciFilter>()?)
} else {
None
};
input.parse::<Token![=>]>()?;
let body: RicciGroup = input.parse()?;
Ok(Self {
index_declaration,
alias_declarations,
filter,
body,
})
}
}
impl Display for RicciSegment {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RicciSegment::TensorCall {
tensor_name,
indexers,
} => {
write!(f, " {} ⟦", tensor_name)?;
let mut first = true;
for indexer in indexers {
if !first {
f.write_str(" ,")?;
}
first = false;
indexer.fmt(f)?;
}
write!(f, " ⟧")
}
RicciSegment::SubLambda(lambda) => lambda.fmt(f),
RicciSegment::SubGroup { delimiter, group } => match delimiter {
Delimiter::Parenthesis => write!(f, " ({})", group),
Delimiter::Brace => write!(f, " {{{}}}", group),
Delimiter::Bracket => write!(f, " [{}]", group),
Delimiter::None => group.fmt(f),
},
RicciSegment::Token(token) => {
write!(f, " {}", token)
}
}
}
}
impl Display for RicciGroup {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
for segment in &self.segments {
segment.fmt(f)?;
}
Ok(())
}
}
impl Display for RicciFilter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "?")?;
for segment in &self.segments {
segment.fmt(f)?;
}
Ok(())
}
}
impl Display for RicciLambda {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{} ", self.index_declaration)?;
for alias_declaration in &self.alias_declarations {
write!(f, "{} ", alias_declaration)?;
}
if let Some(filter) = &self.filter {
write!(f, "{} ", filter)?;
}
write!(f, "▸{}", self.body)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::model::lambda::{RicciGroup, RicciLambda};
use quote::quote;
#[test]
fn parse_lambda() {
let tokens = quote!(for i => a[i] + 3);
assert_eq!(
asserts::parse_and_display::<RicciLambda>(tokens),
"∀ i ▸ a ⟦ i ⟧ + 3"
);
let tokens = quote!(for i j let k = sort[j] => a[i, j] + 4 * b[k]);
assert_eq!(
asserts::parse_and_display::<RicciLambda>(tokens),
"∀ i j ∙ k ≔ sort ⦇ j ⦈ ▸ a ⟦ i , j ⟧ + 4 * b ⟦ k ⟧"
);
}
#[test]
fn parse_complex() {
let tokens = quote!(for i => a[i] + sum(for j => b[i, j]));
assert_eq!(
asserts::parse_and_display::<RicciLambda>(tokens),
"∀ i ▸ a ⟦ i ⟧ + sum (∀ j ▸ b ⟦ i , j ⟧)"
);
let tokens = quote!(for i => a[i] + sum(for j => b[every3[i], j]));
assert_eq!(
asserts::parse_and_display::<RicciLambda>(tokens),
"∀ i ▸ a ⟦ i ⟧ + sum (∀ j ▸ b ⟦ every3 ⦇ i ⦈ , j ⟧)"
);
let tokens = quote!(3.5 * median(for i => a[i] + sum(for j => b[every3[i], j])));
assert_eq!(
asserts::parse_and_display::<RicciGroup>(tokens),
" 3.5 * median (∀ i ▸ a ⟦ i ⟧ + sum (∀ j ▸ b ⟦ every3 ⦇ i ⦈ , j ⟧))"
);
let tokens = quote!(3.5 * median(for i => a[i] + sum(for j => if i < j { b[every3[i], j] } else { c[j] + 4.0 })));
assert_eq!(
asserts::parse_and_display::<RicciGroup>(tokens),
" 3.5 * median (∀ i ▸ a ⟦ i ⟧ + sum (∀ j ▸ if i < j { b ⟦ every3 ⦇ i ⦈ , j ⟧} else { c ⟦ j ⟧ + 4.0}))"
);
}
#[test]
fn parse_invalid() {
let tokens = quote!(for i => { while ok { a[i] } });
assert_eq!(
asserts::parse_and_display::<RicciLambda>(tokens),
"Failed to parse type `tensorism_gen::model::lambda::RicciLambda`: Keyword while is illegal."
);
}
#[test]
fn parse_filters() {
let tokens = quote!(for i if b[i] < a[i] && 0.0 <= c[i, i] => a[i]);
assert_eq!(
asserts::parse_and_display::<RicciLambda>(tokens),
"∀ i ? b ⟦ i ⟧ < a ⟦ i ⟧ & & 0.0 < = c ⟦ i , i ⟧ ▸ a ⟦ i ⟧"
);
let tokens = quote!(for i => (for j if tensor2[j] < j => tensor1[i, j]).sum() + i);
assert_eq!(
asserts::parse_and_display::<RicciLambda>(tokens),
"∀ i ▸ (∀ j ? tensor2 ⟦ j ⟧ < j ▸ tensor1 ⟦ i , j ⟧) . sum () + i"
);
}
}