Repository: HigherOrderCO/Bend
File: src/term/transform/linearize.rs (Rust)

Code before:
      count_var_uses_in_term(snd, uses);
    }
    Term::Match { scrutinee, arms } => {
      count_var_uses_in_term(scrutinee, uses);
      for (rule, term) in arms {
        if let Pattern::Num(MatchNum::Succ(Some(nam))) = rule {
          add_var(nam.as_ref(), uses)
        }

        count_var_uses_in_term(term, uses);
      }
    }
    Term::List { .. } => unreachable!("Should have been desugared already"),
    Term::Lnk { .. } | Term::Ref { .. } | Term::Num { .. } | Term::Str { .. } | Term::Era => (),
  }
}

/// Var-declaring terms
fn term_with_bind_to_affine(
  term: &mut Term,
  nam: &mut Option<Name>,
  var_uses: &mut HashMap<Name, Val>,
  let_bodies: &mut HashMap<Name, Term>,
) {
  if let Some(name) = nam {
    if let Some(_) = var_uses.get(name).copied() {
      term_to_affine(term, var_uses, let_bodies);
      let uses = actual_var_use(var_uses, name);
      duplicate_lam(nam, term, uses);
      return;
    }
  }

  term_to_affine(term, var_uses, let_bodies);
}

fn term_to_affine(term: &mut Term, var_uses: &mut HashMap<Name, Val>, let_bodies: &mut HashMap<Name, Term>) {
  match term {
    Term::Lam { nam, bod, .. } => term_with_bind_to_affine(bod, nam, var_uses, let_bodies),

    Term::Let { pat: Pattern::Var(Some(nam)), val, nxt } => {
      let uses = var_uses[nam];
      match uses {
        0 => {
          if val.has_unscoped() {
            term_to_affine(val, var_uses, let_bodies);
            term_to_affine(nxt, var_uses, let_bodies);

            let Term::Let { val, nxt, .. } = std::mem::take(term) else { unreachable!() };
            *term = Term::Let { pat: Pattern::Var(None), val, nxt };


Review comment (suggestion):
If we don't check the number of uses anymore here, you can change it to `var_uses.contains_key(name)`

Code after:
      count_var_uses_in_term(snd, uses);
    }
    Term::Match { scrutinee, arms } => {
      count_var_uses_in_term(scrutinee, uses);
      for (rule, term) in arms {
        if let Pattern::Num(MatchNum::Succ(Some(nam))) = rule {
          add_var(nam.as_ref(), uses)
        }

        count_var_uses_in_term(term, uses);
      }
    }
    Term::List { .. } => unreachable!("Should have been desugared already"),
    Term::Lnk { .. } | Term::Ref { .. } | Term::Num { .. } | Term::Str { .. } | Term::Era => (),
  }
}

/// Var-declaring terms
fn term_with_bind_to_affine(
  term: &mut Term,
  nam: &mut Option<Name>,
  var_uses: &mut HashMap<Name, Val>,
  inst_count: &mut HashMap<Name, Val>,
  let_bodies: &mut HashMap<Name, Term>,
) {
  if let Some(name) = nam {
    if var_uses.contains_key(name) {
      term_to_affine(term, var_uses, inst_count, let_bodies);
      let instantiated_count = get_var_uses(Some(name), inst_count);
      duplicate_lam(nam, term, instantiated_count);
      return;
    }
  }

  term_to_affine(term, var_uses, inst_count, let_bodies);
}

fn term_to_affine(
  term: &mut Term,
  var_uses: &mut HashMap<Name, Val>,
  // Count to number of times a `Term::Var { nam }` has been reached without being linearized out
  inst_count: &mut HashMap<Name, Val>,
  let_bodies: &mut HashMap<Name, Term>,
) {
  match term {
    Term::Lam { nam, bod, .. } => term_with_bind_to_affine(bod, nam, var_uses, inst_count, let_bodies),

    Term::Let { pat: Pattern::Var(Some(nam)), val, nxt } => {
      let uses = var_uses[nam];
      match uses {
        0 => {