1 files changed, 95 insertions, 0 deletions
diff --git a/src/test/scala/firrtl/analysis/SymbolTableSpec.scala b/src/test/scala/firrtl/analysis/SymbolTableSpec.scala
new file mode 100644
index 00000000..599b4e52
--- /dev/null
+++ b/src/test/scala/firrtl/analysis/SymbolTableSpec.scala
@@ -0,0 +1,95 @@
+// See LICENSE for license details.
+
+package firrtl.analysis
+
+import firrtl.analyses._
+import firrtl.ir
+import firrtl.options.Dependency
+import org.scalatest.flatspec.AnyFlatSpec
+
+class SymbolTableSpec extends AnyFlatSpec {
+  behavior of "SymbolTable"
+
+  private val src =
+    """circuit m:
+      |  module child:
+      |    input x : UInt<2>
+      |    skip
+      |  module m:
+      |    input clk : Clock
+      |    input x : UInt<1>
+      |    output y : UInt<3>
+      |    wire z : SInt<1>
+      |    node a = cat(asUInt(z), x)
+      |    inst i of child
+      |    reg r: SInt<4>, clk
+      |    mem m:
+      |      data-type => UInt<8>
+      |      depth => 31
+      |      reader => r
+      |      read-latency => 1
+      |      write-latency => 1
+      |      read-under-write => undefined
+      |""".stripMargin
+
+  it should "find all declarations in module m before InferTypes" in {
+    val c = firrtl.Parser.parse(src)
+    val m = c.modules.find(_.name == "m").get
+
+    val syms = SymbolTable.scanModule(m, new LocalSymbolTable with WithMap)
+    assert(syms.size == 8)
+    assert(syms("clk").tpe == ir.ClockType && syms("clk").kind == firrtl.PortKind)
+    assert(syms("x").tpe == ir.UIntType(ir.IntWidth(1)) && syms("x").kind == firrtl.PortKind)
+    assert(syms("y").tpe == ir.UIntType(ir.IntWidth(3)) && syms("y").kind == firrtl.PortKind)
+    assert(syms("z").tpe == ir.SIntType(ir.IntWidth(1)) && syms("z").kind == firrtl.WireKind)
+    // The expression type which determines the node type is only known after InferTypes.
+    assert(syms("a").tpe == ir.UnknownType && syms("a").kind == firrtl.NodeKind)
+    // The type of the instance is unknown because we scanned the module before InferTypes and the table
+    // uses only local information.
+    assert(syms("i").tpe == ir.UnknownType && syms("i").kind == firrtl.InstanceKind)
+    assert(syms("r").tpe == ir.SIntType(ir.IntWidth(4)) && syms("r").kind == firrtl.RegKind)
+    val mType = firrtl.passes.MemPortUtils.memType(
+      // only dataType, depth and reader, writer, readwriter properties affect the data type
+      ir.DefMemory(ir.NoInfo, "???", ir.UIntType(ir.IntWidth(8)), 32, 10, 10, Seq("r"), Seq(), Seq(), ir.ReadUnderWrite.New)
+    )
+    assert(syms("m") .tpe == mType && syms("m").kind == firrtl.MemKind)
+  }
+
+  it should "find all declarations in module m after InferTypes" in {
+    val c = firrtl.Parser.parse(src)
+    val inferTypesCompiler = new firrtl.stage.TransformManager(Seq(Dependency(firrtl.passes.InferTypes)))
+    val inferredC = inferTypesCompiler.execute(firrtl.CircuitState(c, Seq())).circuit
+    val m = inferredC.modules.find(_.name == "m").get
+
+    val syms = SymbolTable.scanModule(m, new LocalSymbolTable with WithMap)
+    // The node type is now known
+    assert(syms("a").tpe == ir.UIntType(ir.IntWidth(2)) && syms("a").kind == firrtl.NodeKind)
+    // The type of the instance is now known because it has been filled in by InferTypes.
+    val iType = ir.BundleType(Seq(ir.Field("x", ir.Flip, ir.UIntType(ir.IntWidth(2)))))
+    assert(syms("i").tpe == iType && syms("i").kind == firrtl.InstanceKind)
+  }
+
+  behavior of "WithSeq"
+
+  it should "preserve declaration order" in {
+    val c = firrtl.Parser.parse(src)
+    val m = c.modules.find(_.name == "m").get
+
+    val syms = SymbolTable.scanModule(m, new LocalSymbolTable with WithSeq)
+    assert(syms.getSymbols.map(_.name) == Seq("clk", "x", "y", "z", "a", "i", "r", "m"))
+  }
+
+  behavior of "ModuleTypesSymbolTable"
+
+  it should "derive the module type from the module types map" in {
+    val c = firrtl.Parser.parse(src)
+    val m = c.modules.find(_.name == "m").get
+
+    val childType = ir.BundleType(Seq(ir.Field("x", ir.Flip, ir.UIntType(ir.IntWidth(2)))))
+    val moduleTypes = Map("child" -> childType)
+
+    val syms = SymbolTable.scanModule(m, new ModuleTypesSymbolTable(moduleTypes) with WithMap)
+    assert(syms.size == 8)
+    assert(syms("i").tpe == childType && syms("i").kind == firrtl.InstanceKind)
+  }
+}