1 files changed, 38 insertions, 5 deletions
diff --git a/src/ir/branch-utils.h b/src/ir/branch-utils.h
index cf21facbc..662f70725 100644
--- a/src/ir/branch-utils.h
+++ b/src/ir/branch-utils.h
@@ -236,6 +236,22 @@ inline NameSet getBranchTargets(Expression* ast) {
   return scanner.targets;
 }
 
+// Get the name of the branch target that is defined in the expression, or an
+// empty name if there is none.
+inline Name getDefinedName(Expression* curr) {
+  Name ret;
+  operateOnScopeNameDefs(curr, [&](Name& name) { ret = name; });
+  return ret;
+}
+
+// Return the value sent by a branch instruction, or nullptr if there is none.
+inline Expression* getSentValue(Expression* curr) {
+  Expression* ret = nullptr;
+  operateOnScopeNameUsesAndSentValues(
+    curr, [&](Name name, Expression* value) { ret = value; });
+  return ret;
+}
+
 // Finds if there are branches targeting a name. Note that since names are
 // unique in our IR, we just need to look for the name, and do not need
 // to analyze scoping.
@@ -352,24 +368,41 @@ public:
   }
 };
 
-// Provides a map of branch target names to the expressions that branch to
-// them.
+// Stores information about branch targets, specifically, finding them by their
+// name, and finding the branches to them.
 struct BranchTargets {
   BranchTargets(Expression* expr) { inner.walk(expr); }
 
-  Expression* getTarget(Name name) { return inner.map[name]; }
+  // Gets the expression that defines this branch target, i.e., where we branch
+  // to if we branch to that name.
+  Expression* getTarget(Name name) { return inner.targets[name]; }
+
+  // Gets the expressions branching to a target.
+  std::unordered_set<Expression*> getBranches(Name name) {
+    auto iter = inner.branches.find(name);
+    if (iter != inner.branches.end()) {
+      return iter->second;
+    }
+    return {};
+  }
 
 private:
   struct Inner : public PostWalker<Inner, UnifiedExpressionVisitor<Inner>> {
     void visitExpression(Expression* curr) {
       operateOnScopeNameDefs(curr, [&](Name name) {
         if (name.is()) {
-          map[name] = curr;
+          targets[name] = curr;
+        }
+      });
+      operateOnScopeNameUses(curr, [&](Name& name) {
+        if (name.is()) {
+          branches[name].insert(curr);
         }
       });
     }
 
-    std::map<Name, Expression*> map;
+    std::map<Name, Expression*> targets;
+    std::map<Name, std::unordered_set<Expression*>> branches;
   } inner;
 };