Repository: HigherOrderCO/Bend
File: src/imp/lift_local_defs.rs (Rust)

Code before:
use indexmap::IndexMap;

use crate::fun::{self, Name, Pattern};

use super::{Definition, Expr, Stmt};

impl Definition {
  pub fn lift_local_defs(&mut self, gen: &mut usize) -> Result<IndexMap<Name, fun::Definition>, String> {
    let mut defs = IndexMap::new();
    self.body.lift_local_defs(&self.name, &mut defs, gen)?;
    Ok(defs)
  }
}

impl Stmt {
  pub fn lift_local_defs(
    &mut self,
    parent: &Name,
    defs: &mut IndexMap<Name, fun::Definition>,
    gen: &mut usize,
  ) -> Result<(), String> {
    match self {
      Stmt::LocalDef { .. } => {
        let Stmt::LocalDef { mut def, mut nxt } = std::mem::take(self) else { unreachable!() };
        let children = def.lift_local_defs(gen)?;
        nxt.lift_local_defs(parent, defs, gen)?;

        let local_name = Name::new(format!("{}__local_{}_{}", parent, gen, def.name));
        *gen += 1;

        let (r#use, mut def, fvs) = gen_use(local_name.clone(), *def, nxt)?;
        *self = r#use;
        apply_closure(&mut def, fvs);

        defs.extend(children);
        defs.insert(def.name.clone(), def);
        Ok(())
      }

      Stmt::Assign { pat: _, val: _, nxt } => {
        if let Some(nxt) = nxt {
          nxt.lift_local_defs(parent, defs, gen)?;
        }
        Ok(())
      }
      Stmt::If { cond: _, then, otherwise, nxt } => {
        then.lift_local_defs(parent, defs, gen)?;
        otherwise.lift_local_defs(parent, defs, gen)?;
        if let Some(nxt) = nxt {
          nxt.lift_local_defs(parent, defs, gen)?;
        }

Review comment (refactor):
These two lines will generate a strange order for the extracted functions. Nested definitions are lifted first, then the ones in the same depth are lifted bottom up.

Also, the generated names are difficult to understand, since local functions with the same name will have similar final names.
It would be a bit easier to read the output if they were named like this:
```py
def Foo():
  def A():
    def X():
  def B():
    def Y():
```
compiling to
```
Foo
Foo__local_0_A
Foo__local_0_A__local_0_X
Foo__local_1_B
Foo__local_1_B__local_0_Y
```

What you need to do is mostly just change the order of operations and reset `gen` for each layer of nested defs.


Code after:
use indexmap::IndexMap;

use crate::fun::{self, Name, Pattern};

use super::{Definition, Expr, Stmt};

impl Definition {
  pub fn lift_local_defs(&mut self, gen: &mut usize) -> Result<IndexMap<Name, fun::Definition>, String> {
    let mut defs = IndexMap::new();
    self.body.lift_local_defs(&self.name, &mut defs, gen)?;
    Ok(defs)
  }
}

impl Stmt {
  pub fn lift_local_defs(
    &mut self,
    parent: &Name,
    defs: &mut IndexMap<Name, fun::Definition>,
    gen: &mut usize,
  ) -> Result<(), String> {
    match self {
      Stmt::LocalDef { .. } => {
        let Stmt::LocalDef { mut def, mut nxt } = std::mem::take(self) else { unreachable!() };
        let local_name = Name::new(format!("{}__local_{}_{}", parent, gen, def.name));
        def.body.lift_local_defs(&local_name, defs, gen)?;
        nxt.lift_local_defs(parent, defs, gen)?;
        *gen += 1;

        let inner_defs =
          defs.keys().filter(|name| name.starts_with(local_name.as_ref())).cloned().collect::<BTreeSet<_>>();
        let (r#use, mut def, fvs) = gen_use(local_name.clone(), *def, nxt, inner_defs)?;
        *self = r#use;
        apply_closure(&mut def, fvs);

        defs.insert(def.name.clone(), def);
        Ok(())
      }

      Stmt::Assign { pat: _, val: _, nxt } => {
        if let Some(nxt) = nxt {
          nxt.lift_local_defs(parent, defs, gen)?;
        }
        Ok(())
      }
      Stmt::If { cond: _, then, otherwise, nxt } => {
        then.lift_local_defs(parent, defs, gen)?;
        otherwise.lift_local_defs(parent, defs, gen)?;
        if let Some(nxt) = nxt {
          nxt.lift_local_defs(parent, defs, gen)?;
        }