From a99469a5df8223b876b33dd72b06709da2cd4437 Mon Sep 17 00:00:00 2001
From: Albert Magyar <albert.magyar@gmail.com>
Date: Thu, 14 May 2020 17:29:47 -0700
Subject: [PATCH] Backport ResolveKinds improvements from #1475 to 1.2.x
 (#1623)

* Use HashMap instead of LinkedHashMap in ResolveKinds

* Backport to 1.2.x

(cherry picked from commit 0f78e2df9d8f03e6847a012873fc2408b62d0692)

Co-authored-by: Albert Magyar <albert.magyar@gmail.com>

* Eliminate unnecessary traversals in ResolveKinds

* Modify for backporting

* Make find_port return Unit and use Foreachers in ResolveKinds

* Modify for backporting

Co-authored-by: Jack Koenig <koenig@sifive.com>
---
 .../scala/firrtl/passes/ResolveKinds.scala    | 45 ++++++++++++++-----
 1 file changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/main/scala/firrtl/passes/ResolveKinds.scala b/src/main/scala/firrtl/passes/ResolveKinds.scala
index fd0df5a460..c7d0cdce5e 100644
--- a/src/main/scala/firrtl/passes/ResolveKinds.scala
+++ b/src/main/scala/firrtl/passes/ResolveKinds.scala
@@ -5,8 +5,42 @@ package firrtl.passes
 import firrtl._
 import firrtl.ir._
 import firrtl.Mappers._
+import firrtl.traversals.Foreachers._
+import firrtl.options.PreservesAll
 
 object ResolveKinds extends Pass {
+
+  private def find_port(kinds: collection.mutable.HashMap[String, Kind])(p: Port): Unit = {
+    kinds(p.name) = PortKind
+  }
+
+  def resolve_expr(kinds: collection.mutable.HashMap[String, Kind])(e: Expression): Expression = e match {
+    case ex: WRef => ex copy (kind = kinds(ex.name))
+    case _ => e map resolve_expr(kinds)
+  }
+
+  def resolve_stmt(kinds: collection.mutable.HashMap[String, Kind])(s: Statement): Statement = {
+    s match {
+      case sx: DefWire => kinds(sx.name) = WireKind
+      case sx: DefNode => kinds(sx.name) = NodeKind
+      case sx: DefRegister => kinds(sx.name) = RegKind
+      case sx: WDefInstance => kinds(sx.name) = InstanceKind
+      case sx: DefMemory => kinds(sx.name) = MemKind
+      case _ =>
+    }
+    s.map(resolve_stmt(kinds))
+     .map(resolve_expr(kinds))
+  }
+
+  def resolve_kinds(m: DefModule): DefModule = {
+    val kinds = new collection.mutable.HashMap[String, Kind]
+    m.foreach(find_port(kinds))
+    m.map(resolve_stmt(kinds))
+  }
+
+  def run(c: Circuit): Circuit =
+    c copy (modules = c.modules map resolve_kinds)
+
   type KindMap = collection.mutable.LinkedHashMap[String, Kind]
 
   def find_port(kinds: KindMap)(p: Port): Port = {
@@ -32,14 +66,5 @@ object ResolveKinds extends Pass {
 
   def resolve_stmt(kinds: KindMap)(s: Statement): Statement =
     s map resolve_stmt(kinds) map resolve_expr(kinds)
-
-  def resolve_kinds(m: DefModule): DefModule = {
-    val kinds = new KindMap
-    (m map find_port(kinds)
-       map find_stmt(kinds)
-       map resolve_stmt(kinds))
-  }
-
-  def run(c: Circuit): Circuit =
-    c copy (modules = c.modules map resolve_kinds)
 }
+