%start Schedule

%avoid_insert "ID" "INT" "STRING"
%expect-unused Unmatched 'UNMATCHED'

%left '\\'
%left '|'
%left '&'
%left '.' '@'

%%

Schedule -> OperationList
  :               { OperationList::NilStmt() }
  | Expr          { OperationList::FinalExpr($1) }
  | Stmt Schedule { OperationList::ConsStmt($1, Box::new($2)) }
  ;

Stmt -> Stmt
  : 'let' 'ID' '=' Expr ';'
      { Stmt::LetStmt { span: $span, var: span_of_tok($2), expr: $4 } }
  | 'let' '(' Ids ')' '=' Expr ';'
      { Stmt::LetsStmt { span: $span, vars: rev($3), expr: $6 } }
  | 'ID' '=' Expr ';'
      { Stmt::AssignStmt { span: $span, var: span_of_tok($1), rhs: $3 } }
  | Expr ';'
      { Stmt::ExprStmt { span: $span, exp: $1 } }
  | 'fixpoint' FixpointLimit '{' Schedule '}'
      { Stmt::Fixpoint { span: $span, limit: $2, body: Box::new($4) } }
  | MacroDecl
      { Stmt::MacroDecl { span: $span, def: $1 } }
  ;

FixpointLimit -> FixpointLimit
  : { FixpointLimit::NoLimit { span: $span } }
  | 'stop_after' 'INT'
      { FixpointLimit::StopAfter { span: $span, limit: span_of_tok($2) } }
  | 'panic_after' 'INT'
      { FixpointLimit::PanicAfter { span: $span, limit: span_of_tok($2) } }
  | 'print_iter'
      { FixpointLimit::PrintIter { span: $span } }
  ;

Expr -> Expr
  : 'ID' Args Selector
      { Expr::Function { span: $span, name: span_of_tok($1), args: $2, selection: $3 } }
  | 'MACRO' Args Selector
      { Expr::Macro { span: $span, name: span_of_tok($1), args: $2, selection: $3 } }
  | 'ID'
      { Expr::Variable { span: span_of_tok($1) } }
  | 'INT'
      { Expr::Integer { span: span_of_tok($1) } }
  | 'true'
      { Expr::Boolean { span: $span, val: true } }
  | 'false'
      { Expr::Boolean { span: $span, val: false } }
  | 'STRING'
      { Expr::String { span: $span } }
  | Expr '.' 'ID'
      { Expr::Field { span: $span, lhs: Box::new($1), field: span_of_tok($3) } }
  | Expr '.' 'INT'
      { Expr::TupleField { span: $span, lhs: Box::new($1), field: span_of_tok($3) } }
  | Expr '@' 'ID'
      { Expr::Field { span: $span, lhs: Box::new($1), field: span_of_tok($3) } }
  | '(' Exprs ')'
      { Expr::Tuple { span: $span, exps: $2 } }
  | '[' Exprs ']'
      { Expr::Tuple { span: $span, exps: $2 } }
  | '{' Schedule '}'
      { Expr::BlockExpr { span: $span, body: Box::new($2) } }
  | '<' Fields '>'
      { Expr::Record { span: $span, fields: rev($2) } }
  | Expr '\\' Expr
      { Expr::SetOp { span: $span, op: SetOp::Difference, lhs: Box::new($1), rhs: Box::new($3) } }
  | Expr '|' Expr
      { Expr::SetOp { span: $span, op: SetOp::Union, lhs: Box::new($1), rhs: Box::new($3) } }
  | Expr '&' Expr
      { Expr::SetOp { span: $span, op: SetOp::Intersection, lhs: Box::new($1), rhs: Box::new($3) } }
  ;

Args -> Vec<Expr>
  :                { vec![] }
  | '[' RExprs ']' { rev($2) }
  ;

Exprs -> Vec<Expr>
  : RExprs  { rev($1) }
  ;

RExprs -> Vec<Expr>
  :                  { vec![] }
  | Expr             { vec![$1] }
  | Expr ',' RExprs  { snoc($1, $3) }
  ;

Fields -> Vec<(Span, Expr)>
  :                           { vec![] }
  | 'ID' '=' Expr             { vec![(span_of_tok($1), $3)] }
  | 'ID' '=' Expr ',' Fields  { snoc((span_of_tok($1), $3), $5) }
  ;

Selector -> Selector
  : '(' '*' ')'
      { Selector::SelectAll { span: $span } }
  | '(' Exprs ')'
      { Selector::SelectExprs { span: $span, exprs: $2 } }
  ;

MacroDecl -> MacroDecl
  : 'macro_keyword' 'MACRO' Params '(' 'ID' ')' MacroDef
      { MacroDecl {
          name: span_of_tok($2),
          params: rev($3),
          selection_name: span_of_tok($5),
          def: Box::new($7),
        }
      }
  ;

Params -> Vec<Span>
  :                  { vec![] }
  | '[' Ids ']'      { $2 }
  ;

Ids -> Vec<Span>
  :               { vec![] }
  | 'ID'          { vec![span_of_tok($1)] }
  | 'ID' ',' Ids  { snoc(span_of_tok($1), $3) }
  ;

MacroDef -> OperationList : '{' Schedule '}' { $2 };

Unmatched -> () : 'UNMATCHED' {};

%%

use cfgrammar::Span;
use lrlex::DefaultLexeme;

fn snoc<T>(x: T, mut xs: Vec<T>) -> Vec<T> {
  xs.push(x);
  xs
}

fn rev<T>(mut xs: Vec<T>) -> Vec<T> {
  xs.reverse();
  xs
}

fn span_of_tok(t : Result<DefaultLexeme, DefaultLexeme>) -> Span {
  t.map_err(|_| ()).map(|l| l.span()).unwrap()
}

pub enum OperationList {
  NilStmt(),
  FinalExpr(Expr),
  ConsStmt(Stmt, Box<OperationList>),
}

pub enum Stmt {
  LetStmt    { span: Span, var: Span, expr: Expr },
  LetsStmt   { span: Span, vars: Vec<Span>, expr: Expr },
  AssignStmt { span: Span, var: Span, rhs: Expr },
  ExprStmt   { span: Span, exp: Expr },
  Fixpoint   { span: Span, limit: FixpointLimit, body: Box<OperationList> },
  MacroDecl  { span: Span, def: MacroDecl },
}

pub enum FixpointLimit {
  NoLimit    { span: Span },
  StopAfter  { span: Span, limit: Span },
  PanicAfter { span: Span, limit: Span },
  PrintIter  { span: Span },
}

#[derive(Copy, Clone, Debug)]
pub enum SetOp {
  Difference,
  Union,
  Intersection,
}

pub enum Expr {
  Function    { span: Span, name: Span, args: Vec<Expr>, selection: Selector },
  Macro       { span: Span, name: Span, args: Vec<Expr>, selection: Selector },
  Variable    { span: Span },
  Integer     { span: Span },
  Boolean     { span: Span, val: bool },
  String      { span: Span },
  Field       { span: Span, lhs: Box<Expr>, field: Span },
  BlockExpr   { span: Span, body: Box<OperationList> },
  Record      { span: Span, fields: Vec<(Span, Expr)> },
  SetOp       { span: Span, op: SetOp, lhs: Box<Expr>, rhs: Box<Expr> },
  Tuple       { span: Span, exps: Vec<Expr> },
  TupleField  { span: Span, lhs: Box<Expr>, field: Span },
}

pub enum Selector {
  SelectAll   { span: Span },
  SelectExprs { span: Span, exprs: Vec<Expr> },
}

pub struct MacroDecl {
  pub name: Span,
  pub params: Vec<Span>,
  pub selection_name: Span,
  pub def: Box<OperationList>,
}