Merge pull request #2674 from chipsalliance/select-helpers

add Select combinators

docs/src/diplomacy/select_tutorial.md

@@ -0,0 +1,144 @@
+# Select Library
+Chisel provides a [Select library](https://github.com/freechipsproject/chisel3/blob/master/src/main/scala/chisel3/aop/Select.scala)
+that provides a set of methods for finding specific sets of hardware nodes in a
+Chisel design after it has been elaborated. However, Diplomacy abstracts away a
+lot of the underlying Chisel hardware, so the low-level Chisel `Select`
+combinators can be fragile and difficult to use in designs using Diplomacy. To
+help with this, Rocket Chip provides its own `Select` library that operates on
+`LazyModule`s and `Node`s instead of `Module`s and `Wire`s.
+
+We will use the following `LazyModule`s in the examples below.
+```scala mdoc
+import chisel3.Bool
+import freechips.rocketchip.aop.Select
+import freechips.rocketchip.config.Parameters
+import freechips.rocketchip.diplomacy.{
+  BundleBridgeSink,
+  BundleBridgeSource,
+  LazyModule,
+  LazyModuleImp,
+  SimpleLazyModule
+}
+
+class Top(implicit p: Parameters) extends LazyModule {
+  val a = LazyModule(new A)
+  val foo = LazyModule(new Foo)
+
+  val aInput = BundleBridgeSource[Bool](() => Bool())
+  a.input := aInput
+
+  val aOutput = a.output.makeSink
+
+  val fooInput = BundleBridgeSource[Bool](() => Bool())
+  foo.input := fooInput
+
+  val fooOutput = foo.output.makeSink
+
+  lazy val module = new LazyModuleImp(this) {
+    aInput.makeIO
+    fooOutput.makeIO
+    fooInput.bundle := aOutput.bundle
+  }
+}
+
+class Foo(implicit p: Parameters) extends SimpleLazyModule {
+  val bar = LazyModule(new A)
+
+  val input = bar.input
+  val output = bar.output
+}
+
+class A(implicit p: Parameters) extends LazyModule {
+  val b = LazyModule(new Leaf)
+  val c = LazyModule(new Leaf)
+
+  val input = b.input
+  val output = c.output
+
+  val bOutput = b.output.makeSink
+  val cInput = BundleBridgeSource[Bool](() => Bool())
+  c.input := cInput
+  lazy val module = new LazyModuleImp(this) {
+    cInput.bundle := bOutput.bundle
+  }
+}
+
+class Leaf(implicit p: Parameters) extends LazyModule {
+  val input = BundleBridgeSink[Bool]()
+  val output = BundleBridgeSource[Bool](() => Bool())
+
+  lazy val module = new LazyModuleImp(this) {
+    output.bundle := input.bundle
+  }
+}
+
+val top = LazyModule(new Top()(Parameters.empty))
+```
+
+```scala mdoc:invisible
+// Select methods can only be used after instantiation
+chisel3.stage.ChiselStage.elaborate(top.module)
+```
+
+The instance `top` has the following hierarchy:
+```
+        top:Top
+        /      \
+      a:A      foo:Foo
+     /  \        \
+b:Leaf  c:Leaf   bar:A
+                  /  \
+              b:Leaf  c:Leaf
+```
+
+## Examples
+`Select.collectDeep` takes a `LazyModule` and a partial function. It
+recursively applies the partial to the module and all of its children. The
+following example uses `Select.collectDeep` to collect all `Leaf` modules in
+`top`.
+```scala mdoc
+Select.collectDeep(top) {
+  case l: Leaf => l.pathName
+}
+```
+
+`Select.filterCollectDeep` takes a `LazyModule`, a filter function, and a
+partial function. It recursively applies the partial to the module and all of
+its children if the filter function returns true. When the filter function
+returns `false`, the recursion stops.  The following example uses
+`Select.filterCollectDeep` to collect all `Leaf` modules in `top` that do not
+belong to the `foo:Foo` subhierarchy.
+```scala mdoc
+Select.filterCollectDeep(top) {
+  case _: Foo => false
+  case _ => true
+} {
+  case l: Leaf => l.pathName
+}
+```
+
+`LazyModule`s have a `getNodes` method that return all the nodes instantiated
+within that module. This can be combined with the `Select.collectInwardEdges`
+to select `LazyModule`s based on their connectivity.
+`Select.collectInwardEdges` takes a `BaseNode` and a partial function. It
+applies the function to all the `InwardEdge`s of the node. There is also a
+`Select.collectOutwardEdges` that does the same for outward edges. The example
+below uses `Select.filterCollectDeep`, `LazyModule.getNodes`, and
+`Select.collectInwardEdges` to find all instances of `Leaf` modules in `top`
+that are connected to an `A` module and do not belong to the `foo:Foo`
+subhierarchy.
+```scala mdoc
+Select.filterCollectDeep (top) {
+  case _: Foo => false
+  case _ => true
+} {
+  case a: A =>
+    // Function.unlift used to convert `InwardEdge => Option[String]` to `PartialFunction[InwardEdge, String]
+    a.getNodes.flatMap(Select.collectInwardEdges(_)(Function.unlift { edge =>
+      edge.node.lazyModule match {
+        case l: Leaf => Some(l.pathName)
+        case _ => None
+      }
+    }))
+}
+```

src/main/scala/aop/Select.scala

@@ -0,0 +1,121 @@
+// See LICENSE.SiFive for license details.
+
+package freechips.rocketchip.aop
+
+import chisel3.Data
+import freechips.rocketchip.config.Parameters
+import freechips.rocketchip.diplomacy.{
+  AnyMixedNode,
+  BaseNode,
+  InwardNode,
+  LazyModule,
+  MixedNode,
+  OutwardNode,
+}
+
+/** Combinators for finding specific sets of [[LazyModule]]s/[[Node]]s.
+  *
+  * These can be used for e.g. finding specific TLBundles in a design and
+  * placing monitors or annotating metadata.
+  */
+object Select {
+
+  /** Contains information about an inward edge of a node
+    */
+  case class InwardEdge[Bundle <: Data, EdgeInParams](
+    params: Parameters,
+    bundle: Bundle,
+    edge: EdgeInParams,
+    node: OutwardNode[_, _, Bundle],
+  )
+
+  /** Contains information about an outward edge of a node
+    */
+  case class OutwardEdge[Bundle <: Data, EdgeOutParams](
+    params: Parameters,
+    bundle: Bundle,
+    edge: EdgeOutParams,
+    node: InwardNode[_, _, Bundle],
+  )
+
+  /** Collects the [[InwardEdge]]s of a node. Defined as a separate method so
+    * that the bundle/edge types can be set properly
+    */
+  private def getInwardEdges[BI <: Data, EI](node: MixedNode[_, _, EI, BI, _, _, _, _ <: Data]): Iterable[InwardEdge[BI, EI]] = {
+    node.iPorts.zip(node.in).map {
+      case ((_, node, params, _), (bundle, edge)) =>
+        InwardEdge(params, bundle, edge, node)
+    }
+  }
+
+  /** Applies the collect function to each [[InwardEdge]] of a node
+    */
+  def collectInwardEdges[T](node: BaseNode)(collect: PartialFunction[InwardEdge[_ <: Data, _], T]): Iterable[T] = {
+    node match {
+      case node: AnyMixedNode => getInwardEdges(node).collect(collect)
+      case _ => Seq.empty
+    }
+  }
+
+  /** Collects the [[OutwardEdge]]s of a node. Defined as a separate method so
+    * that the bundle/edge types can be set properly
+    */
+  private def getOutwardEdges[BO <: Data, EO](node: MixedNode[_, _, _, _ <: Data, _, _, EO, BO]): Iterable[OutwardEdge[BO, EO]] = {
+    node.oPorts.zip(node.out).map {
+      case ((_, node, params, _), (bundle, edge)) =>
+        OutwardEdge(params, bundle, edge, node)
+    }
+  }
+
+  /** Applies the collect function to each [[OutardEdge]] of a node
+    */
+  def collectOutwardEdges[T](node: BaseNode)(collect: PartialFunction[OutwardEdge[_ <: Data, _], T]): Iterable[T] = {
+    node match {
+      case node: AnyMixedNode => getOutwardEdges(node).collect(collect)
+      case _ => Seq.empty
+    }
+  }
+
+  /** Applies the collect function to a [[LazyModule]] and recursively to all
+    * of its children.
+    */
+  def collectDeep[T](lmod: LazyModule)(collect: PartialFunction[LazyModule, T]): Iterable[T] = {
+    collect.lift(lmod) ++
+    lmod.getChildren.flatMap { child =>
+      collectDeep(child)(collect)
+    }
+  }
+
+  /** Applies the collect function to a [[LazyModule]] and its children if the
+    * filter function returns true. Stops recursing when the filter function
+    * returns false. e.g.
+    * for this hierarchy
+    *     A
+    *    / \
+    *   B   C
+    *  / \   \
+    * D   E   F
+    *
+    * the following select function
+    * {{{
+    * filterCollectDeep(A) {
+    *   case B => false
+    *   case _ => true
+    * } { m =>
+    *   printl(m)
+    * }
+    * }}}
+    *
+    * will only print modules A, C, and F
+    */
+  def filterCollectDeep[T](lmod: LazyModule)(filter: LazyModule => Boolean)(collect: PartialFunction[LazyModule, T]): Iterable[T] = {
+    if (filter(lmod)) {
+      collect.lift(lmod) ++
+      lmod.getChildren.flatMap { child =>
+        filterCollectDeep(child)(filter)(collect)
+      }
+    } else {
+      Iterable.empty
+    }
+  }
+}

src/main/scala/diplomacy/LazyModule.scala

@@ -5,7 +5,7 @@ package freechips.rocketchip.diplomacy
 import Chisel.{defaultCompileOptions => _, _}
 import chisel3.internal.sourceinfo.{SourceInfo, UnlocatableSourceInfo}
 import chisel3.{MultiIOModule, RawModule, Reset, withClockAndReset}
-import chisel3.experimental.{ChiselAnnotation}
+import chisel3.experimental.ChiselAnnotation
 import firrtl.passes.InlineAnnotation
 import freechips.rocketchip.config.Parameters
 import freechips.rocketchip.util.CompileOptions.NotStrictInferReset

src/main/scala/diplomacy/Nodes.scala

@@ -1507,7 +1507,10 @@ class EphemeralNode[D, U, EO, EI, B <: Data](imp: NodeImp[D, U, EO, EI, B])()(im
   override def omitGraphML = true
   override def oForward(x: Int): Option[(Int, OutwardNode[D, U, B])] = Some(iDirectPorts(x) match { case (i, n, _, _) => (i, n) })
   override def iForward(x: Int): Option[(Int, InwardNode[D, U, B])] = Some(oDirectPorts(x) match { case (i, n, _, _) => (i, n) })
-  override protected[diplomacy] def instantiate(): Seq[Dangle] = Nil
+  override protected[diplomacy] def instantiate(): Seq[Dangle] = {
+    instantiated = true
+    Nil
+  }
 }
 
 /** [[MixedNexusNode]] is used when the number of nodes connecting from either side is unknown (e.g. a Crossbar which also is a protocol adapter).

src/main/scala/diplomacy/package.scala

@@ -179,6 +179,7 @@ import scala.language.implicitConversions
 package object diplomacy
 {
   type SimpleNodeHandle[D, U, E, B <: Chisel.Data] = NodeHandle[D, U, E, B, D, U, E, B]
+  type AnyMixedNode = MixedNode[_, _, _, _ <: Data, _, _, _, _ <: Data]
 
   def sourceLine(sourceInfo: SourceInfo, prefix: String = " (", suffix: String = ")") = sourceInfo match {
     case SourceLine(filename, line, col) => s"$prefix$filename:$line:$col$suffix"