Cleanup aspects (#1359)

* Clean up aspects

* Refactored InjectingAspect with InjectorAspect

* Made AspectLibrary work with objects

* Cleaned up code

* Apply suggestions from code review

* Added tests, removed deprecated newInstance call

* Backed out removal of newInstance as exceptions were different

* Removed trailing commas

6 files changed, 126 insertions, 10 deletions

diff --git a/src/main/resources/META-INF/services/firrtl.options.RegisteredLibrary b/src/main/resources/META-INF/services/firrtl.options.RegisteredLibrary
new file mode 100644
index 00000000..ee160057
--- /dev/null
+++ b/src/main/resources/META-INF/services/firrtl.options.RegisteredLibrary
@@ -0,0 +1 @@
+chisel3.aop.AspectLibrary
diff --git a/src/main/scala/chisel3/aop/AspectLibrary.scala b/src/main/scala/chisel3/aop/AspectLibrary.scala
new file mode 100644
index 00000000..e141688e
--- /dev/null
+++ b/src/main/scala/chisel3/aop/AspectLibrary.scala
@@ -0,0 +1,49 @@
+// See LICENSE for license details.
+
+package chisel3.aop
+
+import chisel3.RawModule
+import firrtl.options.{OptionsException, RegisteredLibrary, ShellOption}
+
+/** Enables adding aspects to a design from the commandline, e.g.
+  *  sbt> runMain chisel3.stage.ChiselMain --module <module> --with-aspect <aspect>
+  */
+final class AspectLibrary() extends RegisteredLibrary  {
+  val name = "AspectLibrary"
+
+  import scala.reflect.runtime.universe._
+
+  private def apply(aspectName: String): Aspect[RawModule] = {
+    try {
+      // If a regular class, instantiate, otherwise try as a singleton object
+      try {
+        val x = Class.forName(aspectName).asInstanceOf[Class[_ <: Aspect[RawModule]]]
+        x.newInstance()
+      } catch {
+        case e: InstantiationException =>
+          val rm = runtimeMirror(getClass.getClassLoader)
+          val x = rm.staticModule(aspectName)
+          try {
+            rm.reflectModule(x).instance.asInstanceOf[Aspect[RawModule]]
+          } catch {
+            case _: Exception => throw e
+          }
+      }
+    } catch {
+      case e: ClassNotFoundException =>
+        throw new OptionsException(s"Unable to locate aspect '$aspectName'! (Did you misspell it?)", e)
+      case e: InstantiationException =>
+        throw new OptionsException(
+          s"Unable to create instance of aspect '$aspectName'! (Does this class take parameters?)", e)
+    }
+  }
+
+  val options = Seq(new ShellOption[String](
+    longOption = "with-aspect",
+    toAnnotationSeq = {
+      case aspectName: String => Seq(apply(aspectName))
+    },
+    helpText = "The name/class of an aspect to compile with (must be a class/object without arguments!)",
+    helpValueName = Some("<package>.<aspect>")
+  ))
+}
diff --git a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
index 00a17d86..ec0b5d28 100644
--- a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
+++ b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
@@ -12,21 +12,35 @@ import firrtl.stage.RunFirrtlTransformAnnotation
 import firrtl.{ir, _}
 
 import scala.collection.mutable
-import scala.reflect.runtime.universe.TypeTag
 
 /** Aspect to inject Chisel code into a module of type M
   *
   * @param selectRoots Given top-level module, pick the instances of a module to apply the aspect (root module)
   * @param injection Function to generate Chisel hardware that will be injected to the end of module m
   *                  Signals in m can be referenced and assigned to as if inside m (yes, it is a bit magical)
-  * @param tTag Needed to prevent type-erasure of the top-level module type
   * @tparam T Type of top-level module
   * @tparam M Type of root module (join point)
   */
-case class InjectingAspect[T <: RawModule,
-                           M <: RawModule](selectRoots: T => Iterable[M],
-                                           injection: M => Unit
-                                          )(implicit tTag: TypeTag[T]) extends Aspect[T] {
+case class InjectingAspect[T <: RawModule, M <: RawModule](
+    selectRoots: T => Iterable[M],
+    injection: M => Unit
+) extends InjectorAspect[T, M](
+    selectRoots,
+    injection
+)
+
+/** Extend to inject Chisel code into a module of type M
+  *
+  * @param selectRoots Given top-level module, pick the instances of a module to apply the aspect (root module)
+  * @param injection Function to generate Chisel hardware that will be injected to the end of module m
+  *                  Signals in m can be referenced and assigned to as if inside m (yes, it is a bit magical)
+  * @tparam T Type of top-level module
+  * @tparam M Type of root module (join point)
+  */
+abstract class InjectorAspect[T <: RawModule, M <: RawModule](
+    selectRoots: T => Iterable[M],
+    injection: M => Unit
+) extends Aspect[T] {
   final def toAnnotation(top: T): AnnotationSeq = {
     toAnnotation(selectRoots(top), top.name)
   }
diff --git a/src/main/scala/chisel3/aop/inspecting/InspectingAspect.scala b/src/main/scala/chisel3/aop/inspecting/InspectingAspect.scala
new file mode 100644
index 00000000..faff2817
--- /dev/null
+++ b/src/main/scala/chisel3/aop/inspecting/InspectingAspect.scala
@@ -0,0 +1,27 @@
+// See LICENSE for license details.
+
+package chisel3.aop.inspecting
+
+import chisel3.RawModule
+import chisel3.aop.Aspect
+import firrtl.AnnotationSeq
+
+/** Use for inspecting an elaborated design and printing out results
+  *
+  * @param inspect Given top-level design, print things and return nothing
+  * @tparam T Type of top-level module
+  */
+case class InspectingAspect[T <: RawModule](inspect: T => Unit) extends InspectorAspect[T](inspect)
+
+
+/** Extend to make custom inspections of an elaborated design and printing out results
+  *
+  * @param inspect Given top-level design, print things and return nothing
+  * @tparam T Type of top-level module
+  */
+abstract class InspectorAspect[T <: RawModule](inspect: T => Unit) extends Aspect[T] {
+  override def toAnnotation(top: T): AnnotationSeq = {
+    inspect(top)
+    Nil
+  }
+}
diff --git a/src/test/scala/chiselTests/aop/SelectSpec.scala b/src/test/scala/chiselTests/aop/SelectSpec.scala
index 80ab518f..49341ed5 100644
--- a/src/test/scala/chiselTests/aop/SelectSpec.scala
+++ b/src/test/scala/chiselTests/aop/SelectSpec.scala
@@ -33,7 +33,7 @@ class SelectTester(results: Seq[Int]) extends BasicTester {
   }
 }
 
-case class SelectAspect[T <: RawModule, X](selector: T => Seq[X], desired: T => Seq[X])(implicit tTag: TypeTag[T]) extends Aspect[T] {
+case class SelectAspect[T <: RawModule, X](selector: T => Seq[X], desired: T => Seq[X]) extends Aspect[T] {
   override def toAnnotation(top: T): AnnotationSeq = {
     val results = selector(top)
     val desiredSeq = desired(top)
diff --git a/src/test/scala/chiselTests/stage/ChiselMainSpec.scala b/src/test/scala/chiselTests/stage/ChiselMainSpec.scala
index 7862acd6..2b3b9c2c 100644
--- a/src/test/scala/chiselTests/stage/ChiselMainSpec.scala
+++ b/src/test/scala/chiselTests/stage/ChiselMainSpec.scala
@@ -4,22 +4,37 @@ package chiselTests.stage
 
 import chisel3._
 import chisel3.stage.ChiselMain
-
 import java.io.File
 
+import chisel3.aop.inspecting.{InspectingAspect, InspectorAspect}
 import org.scalatest.{FeatureSpec, GivenWhenThen, Matchers}
 
 object ChiselMainSpec {
 
   /** A module that connects two different types together resulting in an elaboration error */
   class DifferentTypesModule extends RawModule {
-    val in = IO(UInt(1.W))
-    val out = IO(SInt(1.W))
+    val in = IO(Input(UInt(1.W)))
+    val out = IO(Output(SInt(1.W)))
+    out := in
+  }
+
+  /** A module that connects two of the same types together */
+  class SameTypesModule extends MultiIOModule {
+    val in = IO(Input(UInt(1.W)))
+    val out = IO(Output(UInt(1.W)))
     out := in
   }
 
 }
 
+case class TestClassAspect() extends InspectorAspect[RawModule] ({
+  _: RawModule => println("Ran inspectingAspect")
+})
+
+case object TestObjectAspect extends InspectorAspect[RawModule] ({
+  _: RawModule => println("Ran inspectingAspect")
+})
+
 class ChiselMainSpec extends FeatureSpec with GivenWhenThen with Matchers with chiselTests.Utils {
 
   import ChiselMainSpec._
@@ -113,5 +128,15 @@ class ChiselMainSpec extends FeatureSpec with GivenWhenThen with Matchers with c
                      result = 1)
     ).foreach(runStageExpectFiles)
   }
+  feature("Aspect library") {
+    Seq(
+      ChiselMainTest(args = Array( "-X", "high", "--with-aspect", "chiselTests.stage.TestClassAspect" ),
+        generator = Some(classOf[SameTypesModule]),
+        stdout = Some("Ran inspectingAspect")),
+      ChiselMainTest(args = Array( "-X", "high", "--with-aspect", "chiselTests.stage.TestObjectAspect" ),
+        generator = Some(classOf[SameTypesModule]),
+        stdout = Some("Ran inspectingAspect"))
+    ).foreach(runStageExpectFiles)
+  }
 
 }