(* diderot.grm
 *
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *)

%name Diderot;

%tokens
  : KW_bool		("bool")
  | KW_die		("die")
  | KW_else		("else")
  | KW_false		("false")
  | KW_field		("field")
  | KW_identity		("identity")	(* identity matrix *)
  | KW_if		("if")
  | KW_image		("image")
  | KW_input		("input")
  | KW_in		("in")
  | KW_initially	("initially")
  | KW_int		("int")
  | KW_kernel		("kernel")
  | KW_load		("load")
  | KW_new		("new")
  | KW_output		("output")
  | KW_print		("print")
  | KW_real		("real")
  | KW_stabilize	("stabilize")
  | KW_strand		("strand")
  | KW_string		("string")
  | KW_tensor		("tensor")
  | KW_true		("true")
  | KW_update		("update")
  | KW_vec2		("vec2")
  | KW_vec3		("vec3")
  | KW_vec4		("vec4")
  | KW_zeros		("zeros")	(* zero tensor *)
  | OP_eq		("=")
  | OP_pluseq		("+=")
  | OP_minuseq		("-=")
  | OP_stareq		("*=")
  | OP_slasheq		("/=")
  | OP_orelse		("||")
  | OP_andalso		("&&")
  | OP_lt		("<")
  | OP_lte		("<=")
  | OP_eqeq		("==")
  | OP_neq		("!=")
  | OP_gte		(">=")
  | OP_gt		(">")
  | OP_plus		("+")
  | OP_minus		("-")
  | OP_star		("*")
  | OP_convolve		("⊛")
  | OP_dot		("•")
  | OP_cross		("×")
  | OP_outer		("⊗")
  | OP_slash		("/")
  | OP_exp		("^")
  | OP_at		("@")
  | OP_D		("∇")
  | OP_Dotimes		("∇⊗")
  | OP_Dtimes		("∇×")
  | OP_Ddot		("∇•")
  | LP			("(")
  | RP			(")")
  | LB			("[")
  | RB			("]")
  | LCB			("{")
  | RCB			("}")
  | COMMA		(",")
  | SEMI		(";")
  | COLON		(":")
  | HASH		("#")
  | BANG		("!")
  | BAR			("|")
  | DOTDOT		("..")
  | INT of IntInf.int
  | FLOAT of FloatLit.float
  | STRING of string
  | ID of Atom.atom
  ;

%keywords
  KW_bool, KW_die, KW_else, KW_field, KW_identity, KW_if, KW_image, KW_initially,
  KW_int, KW_kernel, KW_load, KW_new, KW_output, KW_print, KW_real, KW_stabilize,
  KW_strand, KW_string, KW_tensor, KW_update, KW_vec2, KW_vec3, KW_vec4, KW_zeros;

%defs (
  structure PT = ParseTree
  structure L = Literal
  structure B = BasisNames

(* apply a mark constructor to a span and a tree *)
  fun mark cons (span : AntlrStreamPos.span, tr) = cons{span = span, tree = tr}

(* specialize mark functions for common node types *)
  val markDecl = mark PT.D_Mark
  fun markTy (_, e as PT.T_Mark _) = e
    | markTy (sp, tr) = mark PT.T_Mark (sp, tr)
  fun markStmt (_, e as PT.S_Mark _) = e
    | markStmt (sp, tr) = mark PT.S_Mark (sp, tr)
  fun markExpr (_, e as PT.E_Mark _) = e
    | markExpr (sp, tr) = mark PT.E_Mark (sp, tr)

  fun mkCondExp cons = let
	fun mk (e, []) = e
	  | mk (e, [e']) = cons(e, e')
	  | mk (e, e'::r) = cons(e, mk(e', r))
	in
	  mk
	end

(* build an application for an infix binary operator *)
  fun mkBinApp (e1, rator, e2) = PT.E_BinOp(e1, rator, e2)

(* construct application expressions for left-associative binary operators *)
  fun mkLBinExp (e, []) = e
    | mkLBinExp (e, (id, e')::r) = mkLBinExp (mkBinApp(e, id, e'), r)

(* construct application expressions for right-associative binary operators *)
  fun mkRBinExp (e, []) = e
    | mkRBinExp (e, [(id, e')]) = mkBinApp(e, id, e')
    | mkRBinExp (e, (id, e')::r) = mkBinApp(e, id, mkRBinExp(e', r))

  fun mkOptExp (_, e, NONE) = e
    | mkOptExp (spn, e, SOME mk) = mk(spn, e)

  fun flatten NONE = []
    | flatten (SOME(x, xs)) = x::xs

);

File
	: GlobalDecl+ StrandDecl CoordinationDecl
		=> (PT.Program{span=FULL_SPAN, tree=GlobalDecl @ [StrandDecl, CoordinationDecl]})
	;

GlobalDecl
	: InputDecl
	| VarDecl
		=> (mark PT.D_Mark (FULL_SPAN, PT.D_Var VarDecl))
	;

(***** Inputs *****)

InputDecl
	: "input" InputType ID ("(" STRING ")")? ( "=" Initializer )? ";"
		=> (markDecl(FULL_SPAN, PT.D_Input(InputType, ID, Option.getOpt(SR1, ""), SR2)))
	;

(***** variable declarations *****)

VarDecl
	: Type ID "=" Expr ";"
		=> (mark PT.VD_Mark (FULL_SPAN, PT.VD_Decl(Type, ID, Expr)))
	;


(***** Strands *****)

StrandDecl
	: "strand" ID "(" Params ")" "{" StrandStateDecl* StrandMethod+ "}"
		=> (markDecl(FULL_SPAN, PT.D_Strand{
		      name = ID, params = Params, state = StrandStateDecl,
		      methods = StrandMethod
		    }))
	;

Params
	: ( Param ( "," Param )* )?
		=> (flatten SR)
	;

Param
	: ValueType ID
		=> (mark PT.P_Mark (FULL_SPAN, PT.P_Param(ValueType, ID)))
	;

StrandStateDecl
	: "output" VarDecl
		=> (true, mark PT.VD_Mark (FULL_SPAN, VarDecl))
	| VarDecl
		=> (false, VarDecl)
	;

StrandMethod
	: MethodId Block
		=> (mark PT.M_Mark (FULL_SPAN, PT.M_Method(MethodId, Block)))
	;

MethodId
	: "update"	=> (StrandUtil.Update)
	| "stabilize"	=> (StrandUtil.Stabilize)
	;

Block
	: "{" Stmt* "}"
		=> (markStmt(FULL_SPAN, PT.S_Block Stmt))
	;

Stmt
	: AtomicStmt
		=> (AtomicStmt)
	| %try "if" "(" Expr ")" Stmt "else" Stmt
		=> (markStmt(FULL_SPAN, PT.S_IfThenElse(Expr, Stmt1, Stmt2)))
	| %try "if" "(" Expr ")" Stmt
		=> (markStmt(FULL_SPAN, PT.S_IfThen(Expr, Stmt)))
	;

AtomicStmt
	: Block
		=> (Block)
	| VarDecl
		=> (PT.S_Decl VarDecl)
	| "stabilize" ";"
		=> (markStmt(FULL_SPAN, PT.S_Stabilize))
	| "die" ";"
		=> (markStmt(FULL_SPAN, PT.S_Die))
	| "new" ID "(" Arguments ")" ";"
		=> (markStmt(FULL_SPAN, PT.S_New(ID, Arguments)))
        | "print" "(" Expr ("," Expr)* ")" ";"
                => (markStmt(FULL_SPAN, PT.S_Print(Expr::SR)))
	| ID "=" Expr ";"
		=> (markStmt(FULL_SPAN, PT.S_Assign(ID, Expr)))
	| ID "+=" Expr ";"
		=> (markStmt(FULL_SPAN, PT.S_OpAssign(ID, B.asgn_add, Expr)))
	| ID "-=" Expr ";"
		=> (markStmt(FULL_SPAN, PT.S_OpAssign(ID, B.asgn_sub, Expr)))
	| ID "*=" Expr ";"
		=> (markStmt(FULL_SPAN, PT.S_OpAssign(ID, B.asgn_mul, Expr)))
	| ID "/=" Expr ";"
		=> (markStmt(FULL_SPAN, PT.S_OpAssign(ID, B.asgn_div, Expr)))
	;

Arguments
	: ( Expr ("," Expr)* )?
		=> (flatten SR)
	;


(***** Coordination declaration *****)

CoordinationDecl
	: "initially" (Array | Collection) ";"
		=> (markDecl(FULL_SPAN, SR))
	;

Array
	: "[" Create "|" Iterations "]"
		=> (PT.D_InitialArray(Create, Iterations))
	;

Create
	: ID "(" Arguments ")"
		=> (mark PT.C_Mark (FULL_SPAN, PT.C_Create(ID, Arguments)))
	;

Iterations
	: Iteration ( "," Iteration )*
		=> (Iteration :: SR)
	;

Iteration
	: ID "in" Expr ".." Expr
		=> (mark PT.I_Mark (FULL_SPAN, PT.I_Range(ID, Expr1, Expr2)))
	;

Collection
	: "{" Create  "|" Iterations "}"
		=> (PT.D_InitialCollection(Create, Iterations))
	;

(***** Types *****)

Type	: "image" "(" Dimension ")" Dimensions
		=> (markTy(FULL_SPAN, PT.T_Image{
			shape = Dimensions, dim = Dimension
		      }))
	| "field" "#" Dimension "(" Dimension ")" Dimensions
		=> (markTy(FULL_SPAN, PT.T_Field{
			diff = Dimension1,
			shape = Dimensions,
			dim = Dimension2
		      }))
	| "kernel" "#" Dimension
		=> (markTy(FULL_SPAN, PT.T_Kernel Dimension))
	| ValueType SeqDimensions
		=> (markTy(FULL_SPAN, SeqDimensions ValueType))
	;

SeqDimensions
        : (* empty *)
                => (fn ty => ty)
        | "{" "}"
                => (fn ty => PT.T_DynSeq ty)
        | "{" Dimension "}" SeqDimensions
                => (fn ty => SeqDimensions(PT.T_Seq(ty, Dimension)))
        ;

InputType
	: "image" "(" Dimension ")" Dimensions
		=> (markTy(FULL_SPAN, PT.T_Image{
			shape = Dimensions, dim = Dimension
		      }))
	| ValueType SeqDimensions
		=> (markTy(FULL_SPAN, SeqDimensions ValueType))
	;

ValueType
	: "tensor" Dimensions
		=> (markTy(FULL_SPAN, PT.T_Tensor Dimensions))
	| "vec2"
		=> (markTy(FULL_SPAN, PT.T_Vec 2))
	| "vec3"
		=> (markTy(FULL_SPAN, PT.T_Vec 3))
	| "vec4"
		=> (markTy(FULL_SPAN, PT.T_Vec 4))
	| "bool"
		=> (markTy(FULL_SPAN, PT.T_Bool))
	| "int"
		=> (markTy(FULL_SPAN, PT.T_Int))
	| "real"
		=> (markTy(FULL_SPAN, PT.T_Real))
	| "string"
		=> (markTy(FULL_SPAN, PT.T_String))
	;

Dimensions
	: "[" ( Dimension ("," Dimension)* )? "]"
		=> (flatten SR)
	;

Dimension
	: INT
		=> (INT)
	;

(***** Expressions *****)

Expr
	: TestExpr ( "if" Expr "else" Expr => (Expr1, Expr2) )?
		=> (case SR
		     of NONE => TestExpr
		      | SOME(e1, e2) => markExpr(FULL_SPAN, PT.E_Cond(TestExpr, e1, e2))
		    (* end case *))
	;

TestExpr
	: AndExpr ( "||" AndExpr )*
		=> (mkCondExp PT.E_OrElse (AndExpr, SR))
	;

AndExpr
	: CmpExpr ( "&&" CmpExpr )*
		=> (mkCondExp PT.E_AndAlso (CmpExpr, SR))
	;

CmpExpr
	: AddExpr ( CmpOp AddExpr )*
		=> (mkLBinExp (AddExpr, SR))
	;

CmpOp
	: "<" => (B.op_lt) | "<=" => (B.op_lte)
	| "==" => (B.op_equ) | "!=" => (B.op_neq)
	| ">=" => (B.op_gte) | ">" => (B.op_gt)
	;

AddExpr
	: MulExpr ( AddOp MulExpr )*
		=> (mkLBinExp (MulExpr, SR))
	;

AddOp
	: "+" => (B.op_add) | "-" => (B.op_sub) | "@" => (B.op_at)
	;

MulExpr
	: PowerExpr ( MulOp PowerExpr )*
		=> (mkLBinExp (PowerExpr, SR))
	;

MulOp
	: "*" => (B.op_mul) | "/" => (B.op_div) | "⊛" => (B.op_convolve)
	| "•" => (B.op_dot) | "×" => (B.op_cross) | "⊗" => (B.op_outer)
	;

PowerExpr
	: PrefixExpr ( "^" PrefixExpr => (B.op_exp, PrefixExpr) )*
		=> (mkLBinExp (PrefixExpr, SR))
	;

PrefixExpr
	: SuffixExpr
		=> (SuffixExpr)
	| "-" PrefixExpr
		=> (markExpr(FULL_SPAN, PT.E_UnaryOp(B.op_neg, PrefixExpr)))
	| "!" PrefixExpr
		=> (markExpr(FULL_SPAN, PT.E_UnaryOp(B.op_not, PrefixExpr)))
	;

SuffixExpr
	: DerivExpr Suffix*
		=> (case Suffix
		     of [] => DerivExpr
		      | ss => markExpr(FULL_SPAN, List.foldl (fn (f, e) => f e) DerivExpr ss)
		    (* end case *))
	;

Suffix
	: "(" Arguments ")"
		=> (fn e => markExpr(FULL_SPAN, PT.E_Apply(e, Arguments)))
	| "[" Indices "]"
		=> (fn e => PT.E_Slice(e, Indices))
	| "{" Expr "}"
		=> (fn e => PT.E_Subscript(e, Expr))
	;

Indices
	: Index ( "," Index )*
		=> (Index :: SR)
	;

Index
	: ":"
		=> (NONE)
	| Expr
		=> (SOME Expr)
	;

DerivExpr
	: AtomExpr
		=> (AtomExpr)
	| "∇" AtomExpr
		=> (markExpr(FULL_SPAN, PT.E_UnaryOp(B.op_D, AtomExpr)))
	| "∇⊗" DerivExpr
		=> (markExpr(FULL_SPAN, PT.E_UnaryOp(B.op_Dotimes, DerivExpr)))
	| "∇×" DerivExpr
		=> (markExpr(FULL_SPAN, PT.E_UnaryOp(B.op_Dtimes, DerivExpr)))
	| "∇•" DerivExpr
		=> (markExpr(FULL_SPAN, PT.E_UnaryOp(B.op_Ddot, DerivExpr)))
	;

AtomExpr
	: ID
		=> (markExpr(FULL_SPAN, PT.E_Var ID))
	| KW_identity "[" Dimension "]"
		=> (markExpr(FULL_SPAN, PT.E_Id Dimension))
	| KW_zeros Dimensions
		=> (markExpr(FULL_SPAN, PT.E_Zero Dimensions))
(* should the following be classified as SuffixExprs? *)
	| "real" "(" Expr ")"
		=> (markExpr(FULL_SPAN, PT.E_Real Expr))
	| LoadExpr
		=> (LoadExpr)
	| "(" Expr ("," Expr)* ")"
		=> (case SR
		     of [] => Expr
		      | l => markExpr(FULL_SPAN, PT.E_Tuple(Expr::l))
		    (* end case *))
	| "{" Expr ("," Expr)* "}"
		=> (markExpr(FULL_SPAN, PT.E_Sequence(Expr::SR)))
	| "[" Expr ("," Expr)* "]"
		=> (markExpr(FULL_SPAN, PT.E_Cons(Expr::SR)))
	| INT
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Int INT)))
	| FLOAT
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Float FLOAT)))
	| STRING
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.String STRING)))
	| "true"
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Bool true)))
	| "false"
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Bool false)))
	| "|" Expr "|"
		=> (markExpr(FULL_SPAN, PT.E_UnaryOp(B.op_norm, Expr)))
	;

(***** Initializers (constant expressions) *****)

Initializer
	: KW_identity "[" Dimension "]"
		=> (markExpr(FULL_SPAN, PT.E_Id Dimension))
	| KW_zeros Dimensions
		=> (markExpr(FULL_SPAN, PT.E_Zero Dimensions))
	| LoadExpr
		=> (LoadExpr)
	| "{" Initializer ("," Initializer)* "}"
		=> (markExpr(FULL_SPAN, PT.E_Sequence(Initializer::SR)))
	| "[" Initializer ("," Initializer)* "]"
		=> (markExpr(FULL_SPAN, PT.E_Cons(Initializer::SR)))
        | "-" INT
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Int(~INT))))
        | "-" FLOAT
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Float(FloatLit.negate FLOAT))))
	| INT
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Int INT)))
	| FLOAT
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Float FLOAT)))
	| STRING
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.String STRING)))
	| "true"
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Bool true)))
	| "false"
		=> (markExpr(FULL_SPAN, PT.E_Lit(L.Bool false)))
	;

LoadExpr
	: KW_image "(" STRING ")"
		=> (markExpr(FULL_SPAN,
		      PT.E_Apply(PT.E_Var B.fn_image, [PT.E_Lit(L.String STRING)])))
	| KW_load "(" STRING ")"
		=> (markExpr(FULL_SPAN,
		      PT.E_Apply(PT.E_Var B.fn_load, [PT.E_Lit(L.String STRING)])))
	;