Traversing a Circuit

Understanding IR node children

Writing a Firrtl pass usually requires writing functions which walk the Firrtl datastructure to either collection information or replace IR nodes with new IR nodes.

The IR datastructure is a tree, where each IR node can have some number of children nodes (which in turn can have more children nodes, etc.). IR nodes without children are called leaves.

Different IR nodes can have different children types. The following table shows the possible children type for each IR node type:

+------------+-----------------------------+
|    Node    |          Children           |
+------------+-----------------------------+
| Circuit    | DefModule                   |
| DefModule  | Port, Statement             |
| Port       | Type, Direction             |
| Statement  | Statement, Expression, Type |
| Expression | Expression, Type            |
| Type       | Type, Width                 |
| Width      |                             |
| Direction  |                             |
+------------+-----------------------------+

The map function

To write a function that traverses a Circuit, we need to first understand the functional programming concept map.

Understanding Seq.map

A Scala sequence of strings, can be represented as a tree with a root node Seq and children nodes "a", "b", and "c":

val s = Seq("a", "b", "c")

    Seq
 /   |   \
"a" "b" "c"

Suppose we define a function f that, given a String argument x, concatenates x with itself:

def f(x: String): String = x + x

We can call s.map to return a new Seq[String] whose children are the result of applying f to every child of s:

val s = Seq("a", "b", "c")
def f(x: String): String = x + x  // repeated declaration for clarity
val t = s.map(f)
println(t) // Seq("aa", "bb", "cc")

Understanding Firrtl's map

We use this "mapping" idea to create our own, custom map methods on IR nodes. Suppose we have a DoPrim expression representing 1 + 1; this can be depicted as a tree of expressions with a root node DoPrim:

        DoPrim
     /          \
UIntValue    UIntValue

If we have a function f that takes an Expression argument and returns a new Expression, we can "map" it onto all children Expression of a given IR node, like our DoPrim. This would return the following new DoPrim, whose children are the result of applying f to every Expression child of DoPrim:

        DoPrim
     /          \
f(UIntValue)    f(UIntValue)

Sometimes IR nodes have children of multiple types. For example, Conditionally has both Expression and Statement children. In this case, the map will only apply its function to the children whose type matches the function's input argument type:

val c = Conditionally(info, e, s1, s2) // e: Expression, s1, s2: Statement, info: FileInfo
def fExp(e: Expression): Expression = ...
def fStmt(s: Statement): Statement = ...
c.map(fExp)  // Conditionally(fExp(e), s1, s2)
c.map(fStmt) // Conditionally(e, fStmt(s1), fStmt(s2))

Side comment: Scala has "infix notation", which allows you to drop the . and parenthesis when calling a function which has one argument. Often, we write these map functions with infix notation:

c map fExp  // equivalent to c.map(fExp)
c map fStmt // equivalent to c.map(fStmt)

Pre-order traversal

Post-order traversal

A simple example