Aspect-Oriented Programming for Chisel (#1077)

Added Aspects to Chisel, enabling a mechanism for dependency injection to hardware modules.

3 files changed, 130 insertions, 0 deletions

diff --git a/src/main/scala/chisel3/aop/injecting/InjectStatement.scala b/src/main/scala/chisel3/aop/injecting/InjectStatement.scala
new file mode 100644
index 00000000..c207454d
--- /dev/null
+++ b/src/main/scala/chisel3/aop/injecting/InjectStatement.scala
@@ -0,0 +1,21 @@
+// See LICENSE for license details.
+
+package chisel3.aop.injecting
+
+import chisel3.stage.phases.AspectPhase
+import firrtl.annotations.{Annotation, ModuleTarget, NoTargetAnnotation, SingleTargetAnnotation}
+
+/** Contains all information needed to inject statements into a module
+  *
+  * Generated when a [[InjectingAspect]] is consumed by a [[AspectPhase]]
+  * Consumed by [[InjectingTransform]]
+  *
+  * @param module Module to inject code into at the end of the module
+  * @param s Statements to inject
+  * @param modules Additional modules that may be instantiated by s
+  * @param annotations Additional annotations that should be passed down compiler
+  */
+case class InjectStatement(module: ModuleTarget, s: firrtl.ir.Statement, modules: Seq[firrtl.ir.DefModule], annotations: Seq[Annotation]) extends SingleTargetAnnotation[ModuleTarget] {
+  val target: ModuleTarget = module
+  override def duplicate(n: ModuleTarget): Annotation = this.copy(module = n)
+}
diff --git a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
new file mode 100644
index 00000000..74cd62f3
--- /dev/null
+++ b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
@@ -0,0 +1,63 @@
+// See LICENSE for license details.
+
+package chisel3.aop.injecting
+
+import chisel3.{Module, ModuleAspect, experimental, withClockAndReset}
+import chisel3.aop._
+import chisel3.experimental.RawModule
+import chisel3.internal.Builder
+import chisel3.internal.firrtl.DefModule
+import chisel3.stage.DesignAnnotation
+import firrtl.annotations.ModuleTarget
+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] {
+  final def toAnnotation(top: T): AnnotationSeq = {
+    toAnnotation(selectRoots(top), top.name)
+  }
+
+  final def toAnnotation(modules: Iterable[M], circuit: String): AnnotationSeq = {
+    RunFirrtlTransformAnnotation(new InjectingTransform) +: modules.map { module =>
+      val (chiselIR, _) = Builder.build(Module(new ModuleAspect(module) {
+        module match {
+          case x: experimental.MultiIOModule => withClockAndReset(x.clock, x.reset) { injection(module) }
+          case x: RawModule => injection(module)
+        }
+      }))
+      val comps = chiselIR.components.map {
+        case x: DefModule if x.name == module.name => x.copy(id = module)
+        case other => other
+      }
+
+      val annotations = chiselIR.annotations.map(_.toFirrtl).filterNot{ a => a.isInstanceOf[DesignAnnotation[_]] }
+
+      val stmts = mutable.ArrayBuffer[ir.Statement]()
+      val modules = Aspect.getFirrtl(chiselIR.copy(components = comps)).modules.flatMap {
+        case m: firrtl.ir.Module if m.name == module.name =>
+          stmts += m.body
+          Nil
+        case other =>
+          Seq(other)
+      }
+
+      InjectStatement(ModuleTarget(circuit, module.name), ir.Block(stmts), modules, annotations)
+    }.toSeq
+  }
+}
+
diff --git a/src/main/scala/chisel3/aop/injecting/InjectingTransform.scala b/src/main/scala/chisel3/aop/injecting/InjectingTransform.scala
new file mode 100644
index 00000000..c65bee38
--- /dev/null
+++ b/src/main/scala/chisel3/aop/injecting/InjectingTransform.scala
@@ -0,0 +1,46 @@
+// See LICENSE for license details.
+
+package chisel3.aop.injecting
+
+import firrtl.{ChirrtlForm, CircuitForm, CircuitState, Transform, ir}
+
+import scala.collection.mutable
+
+/** Appends statements contained in [[InjectStatement]] annotations to the end of their corresponding modules
+  *
+  * Implemented with Chisel Aspects and the [[chisel3.aop.injecting]] library
+  */
+class InjectingTransform extends Transform {
+  override def inputForm: CircuitForm = ChirrtlForm
+  override def outputForm: CircuitForm = ChirrtlForm
+
+  override def execute(state: CircuitState): CircuitState = {
+
+    val addStmtMap = mutable.HashMap[String, Seq[ir.Statement]]()
+    val addModules = mutable.ArrayBuffer[ir.DefModule]()
+
+    // Populate addStmtMap and addModules, return annotations in InjectStatements, and omit InjectStatement annotation
+    val newAnnotations = state.annotations.flatMap {
+      case InjectStatement(mt, s, addedModules, annotations) =>
+        addModules ++= addedModules
+        addStmtMap(mt.module) = s +: addStmtMap.getOrElse(mt.module, Nil)
+        annotations
+      case other => Seq(other)
+    }
+
+    // Append all statements to end of corresponding modules
+    val newModules = state.circuit.modules.map { m: ir.DefModule =>
+      m match {
+        case m: ir.Module if addStmtMap.contains(m.name) =>
+          m.copy(body = ir.Block(m.body +: addStmtMap(m.name)))
+        case m: _root_.firrtl.ir.ExtModule if addStmtMap.contains(m.name) =>
+          ir.Module(m.info, m.name, m.ports, ir.Block(addStmtMap(m.name)))
+        case other: ir.DefModule => other
+      }
+    }
+
+    // Return updated circuit and annotations
+    val newCircuit = state.circuit.copy(modules = newModules ++ addModules)
+    state.copy(annotations = newAnnotations, circuit = newCircuit)
+  }
+}