defpackage firrtl/parser :
   import core
   import verse
   import firrtl/ir2
   import stz/parser
   import firrtl/lexer
   import bigint2
   import firrtl/ir-utils
   import firrtl/chirrtl

;======= Convenience Types ===========
definterface MStat 
defstruct Reader <: MStat :
   value: Symbol
defstruct Writer <: MStat :
   value: Symbol
defstruct ReadWriter <: MStat :
   value: Symbol
defstruct ReadLatency <: MStat :
   value: Int
defstruct WriteLatency <: MStat :
   value: Int
defstruct DataType <: MStat :
   value: Type
defstruct Depth <: MStat :
   value: Int
;======= Convenience Functions ========
defn first-info? (form) -> FileInfo|False :
   match(form) :
      (form:Token) : info(form)
      (form:List) : search(first-info?, form)
      (form) : false

defn first-info (form:List) :
   match(first-info?(form)) :
      (i:FileInfo) : i
      (f:False) : FileInfo()

defn FPE (form, x) :
   throw $ new Exception :
      defmethod print (o:OutputStream, this) :
         print(o, "[~] FIRRTL Parsing Error: ~" << [first-info(form), x])

defn* apply-suffix-ops (x, fs:List) :
   if empty?(fs) : x
   else : apply-suffix-ops(head(fs)(x), tail(fs))

defn parse-stmts (forms:List) :
   val cs = Vector<Stmt>()
   defn* loop (forms:List) :
      match-syntax(forms) :
         () : to-list(cs)
         (?s:#stmt ?rest ...) : (add(cs, s), loop(rest))
         (?rest ...) : FPE(rest, "Expected a statement here.")
   loop(forms)  

;======== Parser Utilities ==============
defn atom? (x) : unwrap-token(x) not-typeof List

defn primop (x:Symbol) : get?(OPERATORS, x, false)
val OPERATORS = HashTable<Symbol, PrimOp>(symbol-hash)
OPERATORS[`add] = ADD-OP
OPERATORS[`sub] = SUB-OP
OPERATORS[`mul] = MUL-OP
OPERATORS[`div] = DIV-OP
OPERATORS[`mod] = MOD-OP
OPERATORS[`quo] = QUO-OP
OPERATORS[`rem] = REM-OP
OPERATORS[`addw] = ADD-WRAP-OP
OPERATORS[`subw] = SUB-WRAP-OP
OPERATORS[`lt] = LESS-OP
OPERATORS[`leq] = LESS-EQ-OP
OPERATORS[`gt] = GREATER-OP
OPERATORS[`geq] = GREATER-EQ-OP
OPERATORS[`eq] = EQUAL-OP
OPERATORS[`neq] = NEQUAL-OP
OPERATORS[`eqv] = EQUIV-OP
OPERATORS[`neqv] = NEQUIV-OP
;OPERATORS[`mux] = MUX-OP
OPERATORS[`pad] = PAD-OP
OPERATORS[`neg] = NEG-OP
OPERATORS[`asUInt] = AS-UINT-OP
OPERATORS[`asSInt] = AS-SINT-OP
OPERATORS[`dshl] = DYN-SHIFT-LEFT-OP
OPERATORS[`dshr] = DYN-SHIFT-RIGHT-OP
OPERATORS[`shl] = SHIFT-LEFT-OP
OPERATORS[`shr] = SHIFT-RIGHT-OP
OPERATORS[`cvt] = CONVERT-OP
OPERATORS[`andr] = BIT-AND-REDUCE-OP
OPERATORS[`orr] = BIT-OR-REDUCE-OP
OPERATORS[`xorr] = BIT-XOR-REDUCE-OP
OPERATORS[`not] = BIT-NOT-OP
OPERATORS[`and] = BIT-AND-OP
OPERATORS[`or] = BIT-OR-OP
OPERATORS[`xor] = BIT-XOR-OP
OPERATORS[`cat] = CONCAT-OP
OPERATORS[`bit] = BIT-SELECT-OP
OPERATORS[`bits] = BITS-SELECT-OP

;======== Parser Rules ==================
defsyntax firrtl :
   ;Useful Atoms
   defrule atoms :
      ;Unconditionally parse next form as identifier.
      id = (?x) when atom?(x) :
         match(unwrap-token(x)) :
            (x:Symbol) : x
            (x) : FPE(form, "Expected an identifier here. Got ~ instead." << [x])

      ;Parses next form if integer literal
      int = (?x) when unwrap-token(x) typeof Int :
         unwrap-token(x)

      string = (?x) when unwrap-token(x) typeof String :
         unwrap-token(x)

      ;Parses next form if long literal
      intorlong = (?x) when unwrap-token(x) typeof Int|Long :
         unwrap-token(x)

      ;Parses next form if symbol
      sym = (?x) when unwrap-token(x) typeof Symbol :
         unwrap-token(x)

   ;Error Handling Productions
   defrule :
      ;Error if not an identifier
      id! = (?x:#id) : x
      id! != () : FPE(form, "Expected an identifier here.")

      ;Error if not =>
      =>! = (=>) : (`=>)
      =>! != () : FPE(form, "Expected => here.")

      ;Error if not a colon
      :! = (:) : (`:)
      :! != () : FPE(form, "Expected a colon here.")

      ;Error if not 'of' keyword
      of! = (of) : `of
      of! != () : FPE(form, "Expected the 'of' keyword here.")

      ;Error if not a =
      =! = (=) : `=
      =! != () : FPE(form, "Expected a '=' here.")

      ;Error if not a single integer
      int$ = (?i:#int ?rest ...) when empty?(rest) : i
      int$ != () : FPE(form, "Expected a single integer literal here.")

      ;Error if not an integer
      int! = (?i:#int) : i
      int! != () : FPE(form, "Expected an integer literal here.")

      ;Error if not a single long integer
      long$ = (?i:#intorlong ?rest ...) when empty?(rest) : 
         match(i) :
            (i:Long) : i
            (i) : to-long(i)
      long$ != () : FPE(form, "Expected a single long integer literal here.")

      ;Error if not a single width
      width$ = (?w:#width ?rest ...) when empty?(rest) : w
      width$ != () : FPE(form, "Expected a single width specifier here.")

      ;Error if not a type
      type! = (?t:#type) : t
      type! != () : FPE(form, "Expected a type here.")

      ;Error if not a vec type
      vectype! = (?t:#type!) :
         FPE(form, "Expected a vector type here.") when t not-typeof VectorType
         t

      ;Error if not an expression
      exp! = (?e:#exp) : e
      exp! != () : FPE(form, "Expected an expression here.")

      ;Error if not a single expression
      exp$ = (?e:#exp ?rest ...) when empty?(rest) : e
      exp$ != () : FPE(form, "Expected a single expression here.")

      ;Error if not a stmt
      stmt! = (?s:#stmt) : s
      stmt! != () : FPE(form, "Expected a statement here.")

      ;Error if not a reference expression
      ref! = (?e:#exp!) :
         FPE(form, "Expected a reference expression here.") when e not-typeof Ref
         e
   
   ;Main Circuit Production
   defrule circuit :
      circuit = (circuit ?name:#id! #:! (?ms:#module ... ?rest ...)) :
         if not empty?(rest) :
            FPE(rest, "Expected a module declaration here.")
         Circuit(first-info(form),ms, name)
      circuit != (circuit) :
         FPE(form, "Invalid syntax for circuit definition.")

   ;Main Module Production
   defrule module :
      ;module = (module ?name:#id! #:! (?ps:#port ... ?rest ...)) :
      ;   InModule(first-info(form), name, ps, Begin(parse-stmts(rest)))
      module = (module ?name:#id! #:! (?ps:#port ... ?cs:#stmt ... ?rest ...)) :
         if not empty?(rest) :
            FPE(rest, "Expected a statement here.")
         InModule(first-info(form),name, ps, Begin(cs))
      module = (extmodule ?name:#id! #:! (?ps:#port ... ?rest ...)) :
         if not empty?(rest) :
            FPE(rest, "Expected a port here.")
         ExModule(first-info(form),name, ps)
      module != (module) :
         FPE(form, "Invalid syntax for module definition.")
      module != (exmodule) :
         FPE(form, "Invalid syntax for exmodule definition.")
   
   defrule port :
      port = (input ?name:#id! #:! ?type:#type!) : Port(first-info(form),name, INPUT, type)
      port = (output ?name:#id! #:! ?type:#type!) : Port(first-info(form),name, OUTPUT, type)
         
   ;Main Type Productions
   defrule type :
      inttype = (UInt<?w:#width$>) : UIntType(w)
      inttype = (UInt) : UIntType(UnknownWidth())
      inttype = (SInt<?w:#width$>) : SIntType(w)
      inttype = (SInt) : SIntType(UnknownWidth())

      clktype = (Clock) : ClockType()

      type = (?t:#typeterm ?ops:#typeop ...) : apply-suffix-ops(t, ops)
      type = (?t:#clktype) : t
      typeop = ((@get ?size:#int$)) : (fn (t) : VectorType(t, size))
      
      typeterm = (?t:#inttype) : t
      typeterm = ({?fs:#field ... ?rest ...}) :
         if not empty?(rest) :
            FPE(rest, "Expected a bundle field declaration here.")
         BundleType(fs)

   defrule field :
      field = (flip ?name:#id! #:! ?type:#type!) : Field(name, REVERSE, type)
      field = (?name:#id #:! ?type:#type!) : Field(name, DEFAULT, type)

   defrule width :
      width = (?x:#int) : IntWidth(x)
      width = (?) : UnknownWidth()

   ;Main Statement Productions
   defrule mstat :
      mstat = (reader #=>! ?name:#id!) : Reader(name)
      mstat = (writer #=>! ?name:#id!) : Writer(name)
      mstat = (read-writer #=>! ?name:#id!) : ReadWriter(name)
      mstat = (read-latency #=>! ?i:#int!) : ReadLatency(i)
      mstat = (write-latency #=>! ?i:#int!) : WriteLatency(i)
      mstat = (data-type #=>! ?t:#type!) : DataType(t)
      mstat = (depth #=>! ?i:#int!) : Depth(i)
   defrule statements :
      stmt = (skip) : Empty()
      stmt = (wire ?name:#id! #:! ?t:#type!) : DefWire(first-info(form),name, t)
      stmt = (reg ?name:#id! #:! ?t:#type! ?clk:#exp! ?reset:#exp! ?init:#exp!) : DefRegister(first-info(form),name, t,clk,reset,init)
      ;stmt = (mem ?name:#id! #:! ?data-type:#type! ?depth:#int ?writers:#id! ... ?wl:#int ?readers:#id! ... ?rl:#int ?readwriters:#id! ...) : 
      ;   DefMemory(first-info(form),name,data-type,depth,wl,rl,readers,writers,readwriters)
      stmt = (cmem ?name:#id! #:! ?t:#vectype! ) : CDefMemory(first-info(form),name,type(t),size(t),false)
      stmt = (smem ?name:#id! #:! ?t:#vectype! ) : CDefMemory(first-info(form),name,type(t),size(t),true)

      stmt = (read mport ?name:#id! #=! ?mem:#id! (@get ?index:#exp!) ?clk:#exp!) :
         CDefMPort(first-info(form),name,UnknownType(),mem,list(index,clk),MRead)
      stmt = (write mport ?name:#id! #=! ?mem:#id! (@get ?index:#exp!) ?clk:#exp!) :
         CDefMPort(first-info(form),name,UnknownType(),mem,list(index,clk),MWrite)
      stmt = (rdwr mport ?name:#id! #=! ?mem:#id! (@get ?index:#exp!) ?clk:#exp!) :
         CDefMPort(first-info(form),name,UnknownType(),mem,list(index,clk),MReadWrite)
      stmt = (infer mport ?name:#id! #=! ?mem:#id! (@get ?index:#exp!) ?clk:#exp!) :
         CDefMPort(first-info(form),name,UnknownType(),mem,list(index,clk),MInfer)

      stmt = (mem ?name:#id! #:! (?ms:#mstat ...)) :
         defn grab (f:MStat -> True|False) :
            map(value,to-list(filter(f,ms))) as List
         defn grab1 (f:MStat -> True|False,s:String) :
            val ls = grab(f)
            if length(ls) != 1 : FPE(form,"Must declare one ~." << [s])
            head(ls)
         val readers = grab({_ typeof Reader})
         val writers = grab({_ typeof Writer})
         val readwriters = grab({_ typeof ReadWriter})
         val write-latency = grab1({_ typeof WriteLatency},"write-latency")
         val read-latency = grab1({_ typeof ReadLatency},"read-latency")
         val depth = grab1({_ typeof Depth},"depth")
         val dt = grab1({_ typeof DataType},"data type")
         DefMemory(first-info(form),name,dt,depth,write-latency,read-latency,readers,writers,readwriters)
      stmt = (inst ?name:#id! #of! ?m:#id!) : DefInstance(first-info(form),name,m)
      stmt = (node ?name:#id! #=! ?e:#exp!) : DefNode(first-info(form),name,e)
      stmt = (poison ?name:#id! #:! ?t:#type!) : DefPoison(first-info(form),name, t)
      stmt = (stop(?clk:#exp, ?en:#exp, ?ret:#int)) : Stop(first-info(form),ret,clk,en)
      stmt = (printf(?clk:#exp ?en:#exp ?str:#string ?es:#exp ...)) : Print(first-info(form),str,es,clk,en)
      stmt = (?s:#stmt/when) : s

      stmt = (?x:#exp <= ?y:#exp!) : Connect(first-info(form),x, y) ;>
      stmt = (?x:#exp <- ?y:#exp!) : BulkConnect(first-info(form),x, y);>
      stmt = (?x:#exp is invalid) : IsInvalid(first-info(form),x)

      ;stmt = ((?s:#stmt ?rest ...)) :
      ;   Begin(List(s, parse-stmts(rest)))
      stmt = ((?s:#stmt ?ss:#stmt ... ?rest ...)) :
         if not empty?(rest) :
            FPE(rest, "Expected a statement here.")
         Begin(List(s, ss))
      stmt = (()) :
         Begin(List())
         
   defrule stmt/when :
      stmt/when = (when ?pred:#exp! #:! ?conseq:#stmt! else ?alt:#stmt/when) :
         Conditionally(first-info(form),pred, conseq, alt)
      stmt/when = (when ?pred:#exp! #:! ?conseq:#stmt! else #:! ?alt:#stmt!) :
         Conditionally(first-info(form),pred, conseq, alt)
      stmt/when = (when ?pred:#exp! #:! ?conseq:#stmt!) :
         Conditionally(first-info(form),pred, conseq, Empty())

   ;Main Expressions
   defrule exp :
      ;Suffix Operators
      exp = (?x:#expterm ?ops:#expop ...) : apply-suffix-ops(x, ops)      
      expop = ((@get ?f:#int)) : (fn (x) : SubIndex(x, f, UnknownType()))
      expop = ((@get ?f:#exp!)) : (fn (x) : SubAccess(x, f, UnknownType()))
      expop = (. ?f:#id!) : (fn (x) : SubField(x, f, UnknownType()))
      
      ;Prefix Operators
      expterm = (?t:#inttype(?v:#string)) :
         val b = BigIntLit(v as String)
         match(t) :
            (t:UIntType) : 
               match(width(t)) :
                  (w:IntWidth) :
                     if to-long(req-num-bits(b)) > width(w) :
                        FPE(form, "Width too small for UIntValue.")
                     UIntValue(b, w)
                  (w) :
                     ;UIntValue(b, w)
                     val num-bits = req-num-bits(b)
                     UIntValue(b,IntWidth(max(1,num-bits)))
            (t:SIntType) : 
               match(width(t)) :
                  (w:IntWidth) :
                     if to-long(req-num-bits(b)) > width(w) :
                        FPE(form, "Width too small for SIntValue.")
                     SIntValue(b, w)
                  (w) :
                     SIntValue(b, w)

      expterm = (?t:#inttype(?v:#int$)) :
         match(t) :
            (t:UIntType) : 
               if (v as Int) < 0 : 
                  FPE(form, "UIntValue cannot be negative.")
               match(width(t)) :
                  (w:IntWidth) : 
                     UIntValue(BigIntLit(v as Int,to-int(width(w)) + 1),w)
                  (w) : 
                     val num-bits = req-num-bits(v as Int)
                     UIntValue(BigIntLit(v as Int,num-bits), IntWidth(max(1,num-bits - 1)))
            (t:SIntType) :
               match(width(t)) :
                  (w:IntWidth) : 
                     SIntValue(BigIntLit(v as Int,to-int(width(w))),w)
                  (w) : 
                     val num-bits = req-num-bits(v as Int)
                     SIntValue(BigIntLit(v as Int,num-bits), IntWidth(num-bits))

      expterm = (mux(?cond:#exp ?tval:#exp ?fval:#exp)) :
         Mux(cond,tval,fval,UnknownType())
      expterm = (validif(?cond:#exp ?value:#exp)) :
         ValidIf(cond,value,UnknownType())
      expterm = (?op:#sym(?es:#exp ... ?ints:#int ... ?rest ...)) :
         if not empty?(rest) :
            FPE(rest, "Illegal operands to primitive operator.")
         match(primop(op)) :
            (p:PrimOp) : DoPrim(p, es, ints, UnknownType())
            (p:False) : FPE(form, "Unrecognized primitive operator '~'." << [op])
      expterm = (?op:#sym) :
         Ref(op, UnknownType())

public defn parse-firrtl (forms:List) :
   with-syntax(firrtl) :
      match-syntax(forms) :
         (?c:#circuit) : c
         (_ ...) : FPE(form, "Invalid firrtl circuit.")

public defn parse-firrtl-file (filename:String) :
   parse-firrtl(lex-file(filename))