cps-ir: Merge variables based on set-based liveness and graph coloring.

pkg/compiler/lib/src/tree_ir/optimization/variable_merger.dart

Index: pkg/compiler/lib/src/tree_ir/optimization/variable_merger.dart
diff --git a/pkg/compiler/lib/src/tree_ir/optimization/variable_merger.dart b/pkg/compiler/lib/src/tree_ir/optimization/variable_merger.dart
new file mode 100644
index 0000000000000000000000000000000000000000..7b5d0057a58ac2117b1fb58c9928c5c2533dbc4e
--- /dev/null
+++ b/pkg/compiler/lib/src/tree_ir/optimization/variable_merger.dart
@@ -0,0 +1,518 @@
+library tree_ir.optimization.variable_merger;
+
+import 'optimization.dart' show Pass, PassMixin;
+import '../tree_ir_nodes.dart';
+import '../../elements/elements.dart' show Local;
+
+/// Merges variables based on liveness and source variable information.
+///
+/// This phase cleans up artifacts introduced by the translation through CPS,
+/// where each source variable is translated into several copies. The copies
+/// are merged again when they are not live simultaneously.
+class VariableMerger extends RecursiveVisitor with PassMixin {
+  String get passName => 'Variable merger';
+
+  @override
+  void rewriteExecutableDefinition(ExecutableDefinition node) {
+    visitExecutableDefinition(node);
+  }
+
+  /// Rewrites the given function.
+  /// This is called for the outermost function and inner functions.
+  void rewriteFunction(ExecutableDefinition node) {
+    BlockGraphBuilder builder = new BlockGraphBuilder();
+    builder.visitExecutableDefinition(node);
+    _computeLiveness(builder.blocks);
+    Map<Variable, Variable> subst = _computeRegisterAllocation(builder.blocks);
+    new SubstVariables(subst).visitExecutableDefinition(node);
+  }
+
+  visitFunctionDefinition(FunctionDefinition node) {
+    super.visitFunctionDefinition(node); // Recurse to visit inner functions.
+    rewriteFunction(node);
+  }
+
+  visitFieldDefinition(FieldDefinition node) {
+    super.visitFieldDefinition(node);
+    rewriteFunction(node);
+  }
+
+  visitConstructorDefinition(ConstructorDefinition node) {
+    super.visitConstructorDefinition(node);
+    rewriteFunction(node);
+  }
+}
+
+/// Basic block in a control-flow graph.
+///
+/// Each block consists of a sequence of reads or a sequence of writes.

[Kevin Millikin (Google), 2015/03/26 14:35:54]

Flesh out this comment a bit to say explicitly that these are basic blocks but
not maximal basic blocks.  Also, indicate why we've chosen to do it that way,
which is what the reader will probably wonder.

[asgerf, 2015/03/27 15:18:26]

Changed to basic blocks with interleaved operations.

+class Block {
+  /// List of predecessors in the control-flow graph.s

[Kevin Millikin (Google), 2015/03/26 14:35:53]

There's an extra 's'.

[asgerf, 2015/03/27 15:18:27]

Done.

+  final List<Block> predecessors = <Block>[];
+
+  /// Entry to the catch block for the enclosing try, or `null`.
+  final Block catchBlock;
+
+  /// List of nodes with this block as [catchBlock].
+  final List<Block> catchPredecessors = <Block>[];
+
+  /// True if this is a sequence of read operations, false if write operations.
+  bool isRead = true;
+  bool get isWrite => !isRead;
+
+  /// Variables being read or written in this block.

[Kevin Millikin (Google), 2015/03/26 14:35:53]

I'd mention either "Sequence of variables" or "in this block, in order" or the
like to distinguish it from a set.

[asgerf, 2015/03/27 15:18:27]

Done.

+  final List<Variable> variables = <Variable>[];
+
+  /// Auxilliary fields used by the liveness analysis.

[Kevin Millikin (Google), 2015/03/26 14:35:54]

One 'l' in auxiliary.

[asgerf, 2015/03/27 15:18:27]

Done.

+  bool inWorklist = true;
+  Set<Variable> liveBefore = new Set<Variable>();

[Kevin Millikin (Google), 2015/03/26 14:35:54]

Consider liveIn and liveOut, which are fairly standard.

[asgerf, 2015/03/27 15:18:26]

Done.

+  Set<Variable> liveAfter = new Set<Variable>();
+
+  Block(this.catchBlock) {
+    if (catchBlock != null) {
+      catchBlock.catchPredecessors.add(this);
+    }
+  }
+}
+
+/// Builds a control-flow graph suitable for performing liveness analysis.
+class BlockGraphBuilder extends RecursiveVisitor {

[Kevin Millikin (Google), 2015/03/26 14:35:55]

The blank line after class isn't necessary.

[asgerf, 2015/03/27 15:18:26]

Done.

+
+  Map<Label, Block> jumpTarget = <Label, Block>{};

[Kevin Millikin (Google), 2015/03/26 14:35:53]

I'd make most of these private names, except for the public api ('blocks'). 
Especially state like currentBlock.

[asgerf, 2015/03/27 15:18:27]

Done.

+  Block currentBlock;
+  List<Block> blocks = <Block>[];
+
+  /// Variables with an assignment that should be treated as final.
+  ///
+  /// Such variables cannot be merged with any other variables, so we exclude
+  /// them from the control-flow graph entirely.
+  Set<Variable> finalVariables = new Set<Variable>();
+
+  BlockGraphBuilder() {
+    currentBlock = newBlock();
+  }
+
+  /// Creates a new block with the current exception handler or [catchBlock]
+  /// if provided.
+  Block newBlock({Block catchBlock}) {
+    if (catchBlock == null && currentBlock != null) {
+      catchBlock = currentBlock.catchBlock;
+    }
+    Block block = new Block(catchBlock);
+    blocks.add(block);
+    return block;
+  }
+
+  /// Starts a new branch after the end of [block].

[Kevin Millikin (Google), 2015/03/26 14:35:53]

Not necessary a branch (at least, not a non-trivial one).  I'd just say "Starts
a new block..." in the comment.  I'm not sure I'd change the function's name,
because I can't think of anything much better.

[asgerf, 2015/03/27 15:18:28]

Done.

+  void branchFrom(Block block) {
+    currentBlock = newBlock()..predecessors.add(block);
+  }
+
+  /// Called when reading from [v].
+  ///
+  /// Appends a read operation to the current basic block, or starts a new
+  /// block if the current block is a write block.
+  void read(Variable v) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

Go ahead and spell out 'variable', here and below.

[asgerf, 2015/03/27 15:18:27]

Done.

+    if (v.isCaptured) return;
+    if (finalVariables.contains(v)) return;
+    if (!currentBlock.isRead) {
+      branchFrom(currentBlock);
+      currentBlock.isRead = true;

[Kevin Millikin (Google), 2015/03/26 14:35:53]

true is the default and we do rely on that a bit below, so perhaps it's clearer
to just assert it.

[asgerf, 2015/03/27 15:18:26]

Not relevant anymore.

+    }
+    currentBlock.variables.add(v);
+  }
+
+  /// Called when writing to [v].
+  ///
+  /// Appends a write operation to the current basic block, or starts a new
+  /// block if the current block is a read block.
+  void write(Variable v) {
+    if (v.isCaptured) return;
+    if (finalVariables.contains(v)) return;
+    if (currentBlock.isRead) {
+      if (!currentBlock.variables.isEmpty) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

This needs a short comment to explain that this code is because the default for
blocks is isRead, so we are handling the first variable use in a fresh block.

[asgerf, 2015/03/27 15:18:27]

Not relevant anymore.

+        branchFrom(currentBlock);
+      }
+      currentBlock.isRead = false;
+    }
+    currentBlock.variables.add(v);
+  }
+
+  /// Called to indicate that [v] has a final assignment, and should therefore
+  /// be ignored. Subsequent calls to [read] and [write] will ignore the it.
+  void finalWrite(Variable v) {
+    finalVariables.add(v);
+  }
+
+  visitVariableUse(VariableUse node) {
+    read(node.variable);
+  }
+
+  visitAssign(Assign node) {
+    visitExpression(node.value);
+    write(node.variable);
+    visitStatement(node.next);
+  }
+
+  visitIf(If node) {
+    visitExpression(node.condition);
+    Block afterCondition = currentBlock;
+    branchFrom(afterCondition);
+    visitStatement(node.thenStatement);
+    branchFrom(afterCondition);
+    visitStatement(node.elseStatement);
+  }
+
+  visitLabeledStatement(LabeledStatement node) {
+    Block join = jumpTarget[node.label] = newBlock();
+    visitStatement(node.body); // visitBreak will add predecessors to join.
+    currentBlock = join;
+    visitStatement(node.next);
+  }
+
+  visitBreak(Break node) {
+    jumpTarget[node.target].predecessors.add(currentBlock);
+  }
+
+  visitContinue(Continue node) {
+    jumpTarget[node.target].predecessors.add(currentBlock);
+  }
+
+  visitWhileTrue(WhileTrue node) {
+    Block join = jumpTarget[node.label] = newBlock();
+    join.predecessors.add(currentBlock);
+    currentBlock = join;
+    visitStatement(node.body); // visitContinue will add predecessors to join.
+  }
+
+  visitWhileCondition(WhileCondition node) {
+    Block join = jumpTarget[node.label] = newBlock();
+    join.predecessors.add(currentBlock);
+    currentBlock = join;
+    visitExpression(node.condition);
+    Block afterCondition = currentBlock;
+    branchFrom(afterCondition);
+    visitStatement(node.body); // visitContinue will add predecessors to join.
+    branchFrom(afterCondition);
+    visitStatement(node.next);
+  }
+
+  visitTry(Try node) {
+    Block catchBlock = newBlock();
+    Block tryBlock = newBlock(catchBlock: catchBlock);

[Kevin Millikin (Google), 2015/03/26 14:35:53]

The three lines starting here are branchFrom with a catch block.  I suggest
making branchFrom (optionally) take a catch block.

[asgerf, 2015/03/27 15:18:27]

Done.

+    tryBlock.predecessors.add(currentBlock);
+    currentBlock = tryBlock;
+    visitStatement(node.tryBody);
+    currentBlock = catchBlock;
+    node.catchParameters.forEach(finalWrite);

[Kevin Millikin (Google), 2015/03/26 14:35:53]

Is this necessary or defensive?  It seems like since the catch parameters are
mutable, we could have phi-assignments to them, and we'd want to try to coalesce
them.

[asgerf, 2015/03/27 15:18:27]

Both. Something smarter could be done, but it's not so simple.

A catch parameter can't be hoisted to the top of the function, so if we merge it
with a variable that is used both inside and outside the catch body, the two
uses will be end up referring to different variables. This is the case for both
JS and Dart.

Since we only merge variables that originated from the same source variable,
this seem unlikely to happen, but I'm not convinced that it can't. For example:

try { 
    return foo();
} catch (e) {
    if (e is! SmartException) e = new SmartException(e);
    e.beSmart();
}

Since the code after catch is only reachable from the catch, it's not hard to
imagine that a CPS rewrite might shuffle the code around a little, so it looks
like this:

var e2;
try {
    return foo();
} catch (e) {
    if (e is! SmartException) {
        e2 = new SmartException(e);
    } else {
        e2 = e;
    }
}
e2.beSmart();

We can't merge e and e2 in this case. Not even in JavaScript, where catch blocks
also their own scope.

Since catch parameters are probably almost never assigned to, this seems like an
acceptable loss.

I've clarified with a comment.

+    visitStatement(node.catchBody);
+  }
+
+  visitConditional(Conditional node) {
+    visitExpression(node.condition);
+    // TODO(asgerf): When assignment expressions are added, this is no longer
+    // sound; then we need to handle as a branch.
+    visitExpression(node.thenExpression);
+    visitExpression(node.elseExpression);
+  }
+
+  visitLogicalOperator(LogicalOperator node) {
+    visitExpression(node.left);
+    // TODO(asgerf): When assignment expressions are added, this is no longer
+    // sound; then we need to handle as a branch.
+    visitExpression(node.right);
+  }
+
+  visitFunctionDeclaration(FunctionDeclaration node) {
+    finalWrite(node.variable);
+    visitStatement(node.next);
+    // Do not traverse inner function.
+  }
+
+  visitFunctionExpression(FunctionExpression node) {
+    // Do not traverse inner function.
+  }
+
+  visitFunctionDefinition(FunctionDefinition node) {
+    // Function parameters are treated as write operations at the entry point,
+    // so they can potentially be merged with other copies of the parameter.
+    // Note that function parameters always have distinct source variables,
+    // so we don't risk accidentally merging two parameters.
+    node.parameters.forEach(write);
+    visitStatement(node.body);
+  }
+
+  visitConstructorDefinition(ConstructorDefinition node) {
+    node.parameters.forEach(write);
+    node.initializers.forEach(visitInitializer);
+    visitStatement(node.body);
+  }
+}
+
+/// Computes liveness information of the given control-flow graph.
+///
+/// The results are stored in [Block.liveBefore] and [Block.liveAfter].
+void _computeLiveness(List<Block> blocks) {
+  List<Block> worklist = new List<Block>.from(blocks);

[Kevin Millikin (Google), 2015/03/26 14:35:53]

Comment that blocks are initially in AST order and that visiting them in reverse
with a LIFO stack is probably a good order.

[asgerf, 2015/03/27 15:18:26]

Done.

+  while (!worklist.isEmpty) {
+    Block block = worklist.removeLast();
+    block.inWorklist = false;
+    Set<Variable> live = new Set<Variable>.from(block.liveAfter);

[Kevin Millikin (Google), 2015/03/26 14:35:53]

This is potentially expensive, it's making an iterable from block.liveAfter,
creating a new LinkedHashSet, and then adding each element from the iterable to
the new set.  Comment that a specialized set implementation would be faster to
copy.

[asgerf, 2015/03/27 15:18:27]

Done.

+    if (block.isRead) {
+      // Reading a variable makes it live before that point.
+      live.addAll(block.variables);

[Kevin Millikin (Google), 2015/03/26 14:35:54]

I think it is worth pointing out that our blocks are either reads or writes, so
in terms of liveness analysis, they either have an empty kill set (for reads) or
an empty gen set, and that the list of variables in the block is the gen set
(for reads) or the kill set.

[asgerf, 2015/03/27 15:18:28]

Not relevant anymore.

+    } else {
+      // Assigning to a variable makes it dead before that point.
+      // Note that when a variable is live at entry to the current catch block,
+      // it remains live before the assignment.
+      // When a variable becomes live at the catch block, it will be removed
+      // from block.variables, so here we can safely mark them all as dead.
+      live.removeAll(block.variables);
+    }
+
+    // If anything changed, propagate liveness backwards.
+    if (block.liveBefore.length < live.length) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

Comment that this is a monotone analysis: the liveIn and liveOut sets start
empty and only grow, so we can use the length to detect changes.

[asgerf, 2015/03/27 15:18:27]

I agree, but there is now an explicit changed flag.

+      block.liveBefore = live;
+
+      // Propagate live variables to predecessors.
+      for (Block pred in block.predecessors) {

[Kevin Millikin (Google), 2015/03/26 14:35:53]

Spell out predecessor.

[asgerf, 2015/03/27 15:18:27]

Done.

+        int size = pred.liveAfter.length;

[Kevin Millikin (Google), 2015/03/26 14:35:54]

size ==> length, or originalLength or the like.

[asgerf, 2015/03/27 15:18:28]

Done.

+        pred.liveAfter.addAll(live);
+        if (pred.liveAfter.length > size && !pred.inWorklist) {
+          worklist.add(pred);
+          pred.inWorklist = true;
+        }
+      }
+
+      // Propagate live variables to catch predecessors.
+      for (Block pred in block.catchPredecessors) {
+        bool changed = false;
+        int size = pred.liveAfter.length;

[Kevin Millikin (Google), 2015/03/26 14:35:54]

Use 'length' in the name, not 'size'.

[asgerf, 2015/03/27 15:18:27]

Done.

+        pred.liveAfter.addAll(live);
+        if (size < pred.liveAfter.length) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

The analogous comparison in the loop above has size on the right, so it should
be on the right here too.

[asgerf, 2015/03/31 12:14:18]

Done.

+          changed = true;
+        }
+        if (pred.isWrite) {
+          // Remove assignments to variables that are live in the catch block.
+          size = pred.variables.length;
+          pred.variables.removeWhere(block.liveBefore.contains);

[Kevin Millikin (Google), 2015/03/26 14:35:55]

block.liveBefore is the same as live, isn't it?  I think you should just use one
or the other and not mix them.

There should also be a comment here that points out that we are removing
variables from a write block, that is from a kill set, so the liveIn set can
only grow (the analysis is still monotone).

[asgerf, 2015/03/27 15:18:27]

Done.

+          if (pred.variables.length < size) {
+            changed = true;
+          }
+        }
+        if (!pred.inWorklist && changed) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

Probably cheaper to check changed before !pred.inWorklist, and it will parallel
the analogous check in the loop above.

[asgerf, 2015/03/27 15:18:28]

Done.

+          worklist.add(pred);
+          pred.inWorklist = true;
+        }
+      }
+    }
+  }
+}
+
+/// Based on liveness information, computes a map of variable substitutions to
+/// merge variables.
+///
+/// Constructs a register interference graph. This is an undirected graph of
+/// variables, with an edge between two variables if they cannot be merged
+/// (because they are live simultaneously).
+///
+/// We then compute a graph coloring, where the color of a node denotes which
+/// variable it will be substituted by.
+///
+/// We never merge variables that originated from distinct source variables,
+/// so we build a separate register interference graph for each source variable.
+Map<Variable, Variable> _computeRegisterAllocation(List<Block> blocks) {
+  Map<Variable, Set<Variable>> edges = new Map<Variable, Set<Variable>>();

[Kevin Millikin (Google), 2015/03/26 14:35:55]

edges ==> interferences

[asgerf, 2015/03/27 15:18:28]

Done.

+
+  // At the assignment to a variable x, add an edge to every variable that is
+  // live after the assignment (if it came from the same source variable).
+  for (Block block in blocks) {
+    if (block.isWrite) {
+      // Group the liveAfter set by source variable.
+      Map<Local, List<Variable>> liveAfter = <Local, List<Variable>>{};
+      for (Variable x in block.liveAfter) {

[Kevin Millikin (Google), 2015/03/26 14:35:53]

x ==> variable

[asgerf, 2015/03/27 15:18:27]

Done.

+        liveAfter.putIfAbsent(x.element, () => <Variable>[]).add(x);
+        edges.putIfAbsent(x, () => new Set<Variable>());
+      }
+      // Add edges for each variable being assigned here.
+      for (Variable x in block.variables.reversed) {
+        edges.putIfAbsent(x, () => new Set<Variable>());
+        List<Variable> live = liveAfter[x.element];
+        if (live != null) {
+          live.remove(x); // Hide from earlier assignments in the block.
+          for (Variable y in live) {
+            edges[x].add(y);
+            edges[y].add(x);
+          }
+        }
+      }
+    }
+  }
+
+  // Sort the variables by descending degree.
+  // The most constrained variables will be assigned a color first.
+  List<Variable> variables = edges.keys.toList();
+  variables.sort((x, y) => edges[y].length - edges[x].length);
+
+  Map<Local, List<Variable>> registers = <Local, List<Variable>>{};
+  Map<Variable, Variable> subst = <Variable, Variable>{};
+
+  for (Variable v1 in variables) {
+    List<Variable> register =
+        registers.putIfAbsent(v1.element, () => <Variable>[]);
+
+    // Find an unused color.
+    Set<Variable> potential = new Set<Variable>.from(register);
+    for (Variable v2 in edges[v1]) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

I don't have a feel for how big the sets of interferences are and so how often
we have to choose new colors.

Does it make sense to have an early exit like:

Set<Variable> potential = new Set<Variable>.from(register);
if (potential.isNotEmpty) {
  for (Variable v2 in edges[v1]) {
    v2 = subst[v2];
    if (v2 != null) {
      potential.remove(v2);
    }
    if (potential.isEmpty) break;
  }
}

[asgerf, 2015/03/27 15:18:27]

When compiling swarm, it's empty about 90% of the time. So it might pay off to
treat that case specially.

I rewrote it a little bit to also avoid building a set from the empty list.

+      Variable v2subst = subst[v2];
+      if (v2subst != null) {
+        potential.remove(v2subst);
+      }
+    }
+    if (potential.isEmpty) {
+      // If no free color was found, add this variable as a new color.
+      register.add(v1);
+      subst[v1] = v1;
+    } else {
+      subst[v1] = potential.first;
+    }
+  }
+
+  return subst;
+}
+
+/// Performs variable substitution and removes redundant assignments.
+class SubstVariables extends RecursiveVisitor {

[Kevin Millikin (Google), 2015/03/26 14:35:53]

I'm not sure what Subst is supposed to be.  The verb 'substitute'?  That's a
good sign we should spell it out.

[asgerf, 2015/03/27 15:18:26]

Done.

+
+  Map<Variable, Variable> subst;
+
+  SubstVariables(this.subst);
+
+  Variable replaceRead(Variable v) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

v ==> variable, w ==> other.

[asgerf, 2015/03/27 15:18:27]

Done.

+    Variable w = subst[v];
+    if (w == null) return v;

[Kevin Millikin (Google), 2015/03/26 14:35:54]

This is the case for final assignments, or does it arise some other way?  It
needs a comment.

[asgerf, 2015/03/27 15:18:26]

Done.

+    w.readCount++;
+    v.readCount--;
+    return w;
+  }
+
+  Variable replaceWrite(Variable v) {
+    Variable w = subst[v];
+    if (w == null) return v;
+    w.writeCount++;
+    v.writeCount--;
+    return w;
+  }
+
+  void replaceParameters(List<Variable> parameters) {
+    for (int i=0; i < parameters.length; i++) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

'i = 0'.

[asgerf, 2015/03/27 15:18:27]

Done.

+      parameters[i] = replaceWrite(parameters[i]);
+    }
+  }
+
+  visitVariableUse(VariableUse node) {
+    node.variable = replaceRead(node.variable);
+  }
+
+  visitFunctionDefinition(FunctionDefinition node) {
+    replaceParameters(node.parameters);
+    node.body = visitStatement(node.body);
+  }
+
+  visitConstructorDefinition(ConstructorDefinition node) {
+    replaceParameters(node.parameters);
+    node.initializers.forEach(visitInitializer);
+    node.body = visitStatement(node.body);
+  }
+
+  visitFieldInitializer(FieldInitializer node) {
+    node.body = visitStatement(node.body);
+  }
+
+  visitSuperInitializer(SuperInitializer node) {
+    for (int i=0; i<node.arguments.length; i++) {

[Kevin Millikin (Google), 2015/03/26 14:35:54]

'i = 0', 'i < node.arguments.length'.

And '++i' :)

[asgerf, 2015/03/27 15:18:26]

Done.

+      node.arguments[i] = visitStatement(node.arguments[i]);
+    }
+  }
+
+  // Statement visitors should return the transformed statement so we
+  // can remove redundant assignments.
+  Statement visitStatement(Statement node) => super.visitStatement(node);
+
+  Statement visitAssign(Assign node) {
+    node.variable = replaceWrite(node.variable);
+
+    visitExpression(node.value);
+    node.next = visitStatement(node.next);
+
+    // Remove assignments of form "x := x"
+    if (node.value is VariableUse) {
+      VariableUse value = node.value;
+      if (value.variable == node.variable) {
+        value.variable.readCount--;
+        node.variable.writeCount--;
+        return node.next;
+      }
+    }
+
+    return node;
+  }
+
+  Statement visitLabeledStatement(LabeledStatement node) {
+    node.body = visitStatement(node.body);
+    node.next = visitStatement(node.next);
+    return node;
+  }
+
+  Statement visitReturn(Return node) {
+    visitExpression(node.value);
+    return node;
+  }
+
+  Statement visitBreak(Break node) => node;

[Kevin Millikin (Google), 2015/03/26 14:35:54]

I usually make all the 'related' method bodies have a parallel style, so

Statement visitBreak(Break node) {
  return node;
}

[asgerf, 2015/03/27 15:18:27]

Done.

+
+  Statement visitContinue(Continue node) => node;
+
+  Statement visitIf(If node) {
+    visitExpression(node.condition);
+    node.thenStatement = visitStatement(node.thenStatement);
+    node.elseStatement = visitStatement(node.elseStatement);
+    return node;
+  }
+
+  Statement visitWhileTrue(WhileTrue node) {
+    node.body = visitStatement(node.body);
+    return node;
+  }
+
+  Statement visitWhileCondition(WhileCondition node) {
+    visitExpression(node.condition);
+    node.body = visitStatement(node.body);
+    node.next = visitStatement(node.next);
+    return node;
+  }
+
+  Statement visitFunctionDeclaration(FunctionDeclaration node) {
+    node.next = visitStatement(node.next);
+    return node;
+  }
+
+  Statement visitExpressionStatement(ExpressionStatement node) {
+    visitExpression(node.expression);
+    node.next = visitStatement(node.next);
+    return node;
+  }
+
+  Statement visitTry(Try node) {
+    node.tryBody = visitStatement(node.tryBody);
+    node.catchBody = visitStatement(node.catchBody);
+    return node;
+  }
+
+  Statement visitSetField(SetField node) {
+    visitExpression(node.object);
+    visitExpression(node.value);
+    node.next = visitStatement(node.next);
+    return node;
+  }
+}