WIP: Added support for FIRRTL 0.2.0 Memories to Scala FIRRTL

4 files changed, 65 insertions, 24 deletions

diff --git a/src/main/antlr4/FIRRTL.g4 b/src/main/antlr4/FIRRTL.g4
index 2765533d..9d186762 100644
--- a/src/main/antlr4/FIRRTL.g4
+++ b/src/main/antlr4/FIRRTL.g4
@@ -47,14 +47,15 @@ block
 stmt
   : 'wire' id ':' type
   | 'reg' id ':' type exp (exp exp)?
-  | 'mem' id ':' '{' 'data-type' '=>' type 
-                     'depth' '=>' IntLit
-                     'read-latency' '=>' IntLit
-                     'write-latency' '=>' IntLit
-                     'read-under-write' '=>' ruw
-                     ('reader' '=>' id)*
-                     ('writer' '=>' id)*
-                     ('readwriter' '=>' id)*
+  | 'mem' id ':' '{' ( 'data-type' '=>' type 
+                     | 'depth' '=>' IntLit
+                     | 'read-latency' '=>' IntLit
+                     | 'write-latency' '=>' IntLit
+                     | 'read-under-write' '=>' ruw
+                     | 'reader' '=>' id
+                     | 'writer' '=>' id
+                     | 'readwriter' '=>' id 
+                     )*
                  '}'
   | 'inst' id 'of' id 
   | 'node' id '=' exp
diff --git a/src/main/scala/firrtl/Translator.scala b/src/main/scala/firrtl/Translator.scala
index 9fe40af8..9f90038c 100644
--- a/src/main/scala/firrtl/Translator.scala
+++ b/src/main/scala/firrtl/Translator.scala
@@ -23,7 +23,7 @@ object Translator
     def countSpaces(s: String): Int = s.prefixLength(_ == ' ')
     def stripComments(s: String): String = s takeWhile (!";".contains(_))
 
-    val Scopers = """\s*(circuit|module|when|else)(.*)""".r
+    val Scopers = """\s*(circuit|module|when|else|mem)(.*)""".r
 
     // Function start
     val it = inputIt.zipWithIndex 
diff --git a/src/main/scala/firrtl/Utils.scala b/src/main/scala/firrtl/Utils.scala
index e0452cfa..46253923 100644
--- a/src/main/scala/firrtl/Utils.scala
+++ b/src/main/scala/firrtl/Utils.scala
@@ -137,16 +137,23 @@ object Utils {
         case r: DefRegister => 
           s"reg ${r.name} : ${r.tpe.serialize}, ${r.clock.serialize}, ${r.reset.serialize}, ${r.init.serialize}"
         case i: DefInstance => s"inst ${i.name} of ${i.module}"
-        case m: DefMemory => s"mem ${m.name} : " + newline + withIndent {
-            s"data-type => ${m.dataType}" + newline +
-            s"depth => ${m.depth}" + newline +
-            s"read-latency => ${m.readLatency}" + newline +
-            s"write-latency => ${m.writeLatency}" + newline +
-            m.readers.map(r => s"reader => ${r}" + newline) +
-            m.writers.map(w => s"writer => ${w}" + newline) +
-            m.readwriters.map(rw => s"readwriter => ${rw}" + newline) +
-            s"read-under-write => undefined" // TODO implement
+        case m: DefMemory => {
+          val str = new StringBuilder(s"mem ${m.name} : " + newline)
+          withIndent {
+            str ++= s"data-type => ${m.dataType.serialize}" + newline +
+              s"depth => ${m.depth}" + newline +
+              s"read-latency => ${m.readLatency}" + newline +
+              s"write-latency => ${m.writeLatency}" + newline +
+              (if (m.readers.nonEmpty) m.readers.map(r => s"reader => ${r}").mkString(newline) + newline
+               else "") +
+              (if (m.writers.nonEmpty) m.writers.map(w => s"writer => ${w}").mkString(newline) + newline
+               else "") + 
+              (if (m.readwriters.nonEmpty) m.readwriters.map(rw => s"readwriter => ${rw}").mkString(newline) + newline
+               else "") +
+              s"read-under-write => undefined"
           }
+          str.result
+        }
         case n: DefNode => s"node ${n.name} = ${n.value.serialize}"
         case c: Connect => s"${c.loc.serialize} <= ${c.exp.serialize}"
         case b: BulkConnect => s"${b.loc.serialize} <- ${b.exp.serialize}"
diff --git a/src/main/scala/firrtl/Visitor.scala b/src/main/scala/firrtl/Visitor.scala
index a1963265..c9d34003 100644
--- a/src/main/scala/firrtl/Visitor.scala
+++ b/src/main/scala/firrtl/Visitor.scala
@@ -1,13 +1,14 @@
 /*
  * TODO
  *  - Implement UBits and SBits
- *  - Support all integer types (not just "h...")
+ *  - Get correct line number for memory field errors
 */
 
 package firrtl
 
 import org.antlr.v4.runtime.tree.AbstractParseTreeVisitor;
 import org.antlr.v4.runtime.ParserRuleContext
+import org.antlr.v4.runtime.tree.ParseTree
 import org.antlr.v4.runtime.tree.ErrorNode
 import org.antlr.v4.runtime.tree.TerminalNode
 import scala.collection.JavaConversions._
@@ -79,21 +80,53 @@ class Visitor(val fullFilename: String) extends FIRRTLBaseVisitor[AST]
 	private def visitBlock[AST](ctx: FIRRTLParser.BlockContext): Stmt = 
     Begin(ctx.stmt.map(visitStmt)) 
 
+  // Memories are fairly complicated to translate thus have a dedicated method
+  private def visitMem[AST](ctx: FIRRTLParser.StmtContext): Stmt = {
+    def parseChildren(children: Seq[ParseTree], map: Map[String, Seq[ParseTree]]): Map[String, Seq[ParseTree]] = {
+      val field = children(0).getText
+      if (field == "}") map
+      else {
+        val newMap = 
+          if (field == "reader" || field == "writer" || field == "readwriter") {
+            val seq = map getOrElse (field, Seq())
+            map + (field -> (seq :+ children(2)))
+          } else { // data-type, depth, read-latency, write-latency, read-under-write
+            if (map.contains(field)) throw new Exception(s"Redefinition of mem field ${field}")
+            else map + (field -> Seq(children(2)))
+          }
+        parseChildren(children.drop(3), newMap) // We consume tokens in groups of three (eg. 'depth' '=>' 5)
+      }
+    }
+
+    val info = getInfo(ctx)
+    // Build map of different Memory fields to their values
+    val map = try {
+      parseChildren(ctx.children.drop(4), Map[String, Seq[ParseTree]]()) // First 4 tokens are 'mem' id ':' '{', skip to fields
+    } catch {
+      case e: Exception => throw new Exception(s"[${info}] ${e.getMessage}") // attach line number
+    }
+    // Each memory field value has been left as ParseTree type, need to convert
+    // TODO Improve? Remove dynamic typecast of data-type
+    DefMemory(info, ctx.id(0).getText, visitType(map("data-type").head.asInstanceOf[FIRRTLParser.TypeContext]), 
+              string2Int(map("depth").head.getText), string2Int(map("write-latency").head.getText), 
+              string2Int(map("read-latency").head.getText), map.getOrElse("reader", Seq()).map(_.getText),
+              map.getOrElse("writer", Seq()).map(_.getText), map.getOrElse("readwriter", Seq()).map(_.getText))
+  }
+
+  // visitStmt
 	private def visitStmt[AST](ctx: FIRRTLParser.StmtContext): Stmt = {
     val info = getInfo(ctx)
 
     ctx.getChild(0).getText match {
-      case "wire" => DefWire(info, ctx.id(0).getText, visitType(ctx.`type`))
+      case "wire" => DefWire(info, ctx.id(0).getText, visitType(ctx.`type`(0)))
       case "reg"  => {
         val name = ctx.id(0).getText
-        val tpe = visitType(ctx.`type`)
+        val tpe = visitType(ctx.`type`(0))
         val (reset, init) = if (ctx.getChildCount > 5) (visitExp(ctx.exp(1)), visitExp(ctx.exp(2))) 
                             else (UIntValue(0, IntWidth(1)), Ref(name, tpe))
         DefRegister(info, name, tpe, visitExp(ctx.exp(0)), reset, init)
       }
-      case "mem" => DefMemory(info, ctx.id(0).getText, visitType(ctx.`type`), string2Int(ctx.IntLit(0).getText),
-                              string2Int(ctx.IntLit(2).getText), string2Int(ctx.IntLit(1).getText),
-                              Seq(), Seq(), Seq()) // TODO implement readers and writers
+      case "mem" => visitMem(ctx)
       case "inst"  => DefInstance(info, ctx.id(0).getText, ctx.id(1).getText)
       case "node" =>  DefNode(info, ctx.id(0).getText, visitExp(ctx.exp(0)))
       case "when" => {