[rust] Recurse into functions in other passes

Updates eliminate_redundant_phis and constant_propagation to recurse into 
function expressions. I also realized there was a bug in EliminateRedundantPhis 
in which we wouldn't traverse into function expressions encountered after 
finding a back edge, so i fixed that logic in both versions.

compiler/forget/crates/fixtures/tests/fixtures/function-expressions.js

@@ -1,5 +1,13 @@
 function Component(props) {
+  const y = 2;
   const foo = function foo(x) {
-    return x + 1;
+    let a = 1;
+    let b;
+    if (a === 1) {
+      b = 5 + 3;
+    } else {
+      b = false;
+    }
+    return x + y + b;
   };
 }

compiler/forget/crates/fixtures/tests/snapshots/fixtures_test__fixtures@function-expressions.js.snap

@@ -5,30 +5,56 @@ input_file: crates/fixtures/tests/fixtures/function-expressions.js
 ---
 Input:
 function Component(props) {
+  const y = 2;
   const foo = function foo(x) {
-    return x + 1;
+    let a = 1;
+    let b;
+    if (a === 1) {
+      b = 5 + 3;
+    } else {
+      b = false;
+    }
+    return x + y + b;
   };
 }
 
 
 Output:
 function Component(
-  unknown props$3,
+  unknown props$6,
 )
 entry bb0
 bb0 (block)
-  [0] #0 = Function @deps[] @context[]:
+  [0] #0 = 2
+  [1] #1 = StoreLocal Const unknown y$7 = unknown #0
+  [2] #2 = Function @deps[] @context[]:
       function foo(
-        unknown x$0,
+        unknown x$8,
       )
       entry bb1
       bb1 (block)
-        [0] #0 = LoadLocal unknown x$0
-        [1] #1 = 1
-        [2] #2 = Binary unknown #0 + unknown #1
-        [3] Return unknown #2
-  [1] #1 = StoreLocal Const unknown foo$4 = unknown #0
-  [2] #2 = <undefined>
-  [3] Return unknown #2
-
-
+        [0] #0 = 1
+        [1] #1 = StoreLocal Let unknown a$9 = unknown #0
+        [2] #2 = DeclareLocal Let unknown b$10
+        [3] #9 = 1
+        [4] #10 = 1
+        [5] #11 = true
+        [6] Goto bb3
+      bb3 (block)
+        predecessors: bb1
+        [7] #3 = 5
+        [8] #4 = 3
+        [9] #5 = 8
+        [10] #6 = StoreLocal Reassign unknown b$11 = unknown #5
+        [11] Goto bb2
+      bb2 (block)
+        predecessors: bb3
+        [12] #12 = LoadLocal unknown x$8
+        [13] #13 = LoadLocal unknown y$0
+        [14] #14 = Binary unknown #12 + unknown #13
+        [15] #15 = 8
+        [16] #16 = Binary unknown #14 + unknown #15
+        [17] Return unknown #16
+  [3] #3 = StoreLocal Const unknown foo$16 = unknown #2
+  [4] #4 = <undefined>
+  [5] Return unknown #4

compiler/forget/crates/hir-optimization/src/constant_propagation.rs

@@ -10,8 +10,16 @@ use hir_ssa::eliminate_redundant_phis;
 
 pub fn constant_propagation<'a>(env: &Environment<'a>, fun: &mut Function<'a>) {
     let mut constants = Constants::new();
+    constant_propagation_impl(env, fun, &mut constants);
+}
+
+fn constant_propagation_impl<'a>(
+    env: &Environment<'a>,
+    fun: &mut Function<'a>,
+    constants: &mut Constants<'a>,
+) {
     loop {
-        let have_terminals_changed = apply_constant_propagation(env, fun, &mut constants);
+        let have_terminals_changed = apply_constant_propagation(env, fun, constants);
         if !have_terminals_changed {
             break;
         }
@@ -162,6 +170,27 @@ fn evaluate_instruction<'a>(
                 constants.insert(value.lvalue.identifier.identifier.id, const_value);
             }
         }
+        InstructionValue::Function(value) => {
+            // TODO: due to the outer fixpoint iteration this could visit the same
+            // function many times. However we only strictly have to visit the function
+            // again if the context variable's constant values have changed since last
+            // time.
+            // Instead, we can:
+            // - Create a filtered Constants instance that extracts just the values for
+            //   the function (using its context variables list)
+            // - Track the last such filtered Constants instance we visited the function
+            //   with. Only visit again if the Constants have changed.
+            let mut inner_constants: Constants<'a> = value
+                .lowered_function
+                .context
+                .iter()
+                .filter_map(|id| {
+                    let value = constants.get(&id.identifier.id);
+                    value.map(|value| (id.identifier.id, value.clone()))
+                })
+                .collect();
+            constant_propagation_impl(env, &mut value.lowered_function, &mut inner_constants);
+        }
         _ => {
             // no-op, not all instructions can be processed
         }

compiler/forget/crates/hir-ssa/src/eliminate_redundant_phis.rs

@@ -1,6 +1,6 @@
 use std::collections::{HashMap, HashSet};
 
-use hir::{BlockId, Environment, Function, Identifier, IdentifierId, HIR};
+use hir::{BlockId, Environment, Function, Identifier, IdentifierId, InstructionValue, HIR};
 use utils::RetainMut;
 
 /// Pass to eliminate redundant phi nodes:
@@ -16,7 +16,7 @@ use utils::RetainMut;
 /// and phis rewrite all their identifiers based on this table. The algorithm loops over the CFG repeatedly
 /// until there are no new rewrites: for a CFG without back-edges it completes in a single pass.
 type Rewrites<'a> = HashMap<IdentifierId, Identifier<'a>>;
-pub fn eliminate_redundant_phis<'a>(_env: &Environment, fun: &mut Function<'a>) {
+pub fn eliminate_redundant_phis<'a>(env: &Environment, fun: &mut Function<'a>) {
     let hir = &mut fun.body;
     let mut rewrites = Rewrites::new();
 
@@ -25,6 +25,7 @@ pub fn eliminate_redundant_phis<'a>(_env: &Environment, fun: &mut Function<'a>)
 
     let mut len;
     loop {
+        let is_first_iteration = !has_back_edge;
         len = rewrites.len();
 
         for (_, block) in hir.blocks.iter_mut() {
@@ -75,6 +76,13 @@ pub fn eliminate_redundant_phis<'a>(_env: &Environment, fun: &mut Function<'a>)
                     rewrite(&rewrites, &mut store.identifier.identifier)
                 });
                 instr.each_identifier_load(|load| rewrite(&rewrites, &mut load.identifier));
+                // Visit function expressions on first iteration of each block to
+                // recursively eliminate any of their redundant phis
+                if is_first_iteration {
+                    if let InstructionValue::Function(fun) = &mut instr.value {
+                        eliminate_redundant_phis(env, &mut fun.lowered_function);
+                    }
+                }
             }
         }
 

compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts

@@ -42,6 +42,7 @@ export function eliminateRedundantPhi(fn: HIRFunction): void {
   // compare to see if any new rewrites were added in that iteration.
   let size = rewrites.size;
   do {
+    const isFirstIteration = !hasBackEdge;
     size = rewrites.size;
     for (const [blockId, block] of ir.blocks) {
       // On the first iteration of the loop check for any back-edges.
@@ -105,7 +106,7 @@ export function eliminateRedundantPhi(fn: HIRFunction): void {
 
         // visit function expressions on first iteration of each block
         if (
-          !hasBackEdge &&
+          isFirstIteration &&
           instr.value.kind === "FunctionExpression" &&
           fn.env.enableOptimizeFunctionExpressions
         ) {