<commit_before>package org.mozartoz.bootcompiler
package transform

import scala.collection.mutable.ListBuffer

import ast._
import symtab._

object Namer extends Transformer with TransformUtils {
  type Env = Map[String, Symbol]

  private var env: Env = _

  override def apply(prog: Program) {
    withEnvironment(prog.builtins.topLevelEnvironment) {
      super.apply(prog)
    }
  }

  private def withEnvironment[A](newEnv: Env)(f: => A) = {
    val savedEnv = env
    env = newEnv
    try f
    finally env = savedEnv
  }

  private def withEnvironmentFromDecls[A](
      decls: List[Variable])(f: => A) = {
    val newEnv = (decls map (decl => decl.name -> decl.symbol))
    withEnvironment(env ++ newEnv)(f)
  }

  override def transformStat(statement: Statement) = statement match {
    case local @ LocalStatement(declarations, body) =>
      val (decls, stats) = extractDecls(declarations)
      val stat = statsAndStatToStat(stats, body)

      withEnvironmentFromDecls(decls) {
        if (decls.isEmpty) transformStat(stat)
        else treeCopy.LocalStatement(local, decls, transformStat(stat))
      }

    case _ =>
      super.transformStat(statement)
  }

  override def transformExpr(expression: Expression) = expression match {
    case local @ LocalExpression(declarations, body) =>
      val (decls, stats) = extractDecls(declarations)
      val expr = statsAndExprToExpr(stats, body)

      withEnvironmentFromDecls(decls) {
        if (decls.isEmpty) transformExpr(expr)
        else treeCopy.LocalExpression(local, decls, transformExpr(expr))
      }

    case proc @ ProcExpression(name, args, body, flags) =>
      val namedFormals = nameFormals(args)

      withEnvironmentFromDecls(namedFormals) {
        treeCopy.ProcExpression(proc,
            name,
            namedFormals,
            transformStat(body),
            flags)
      }

    case fun @ FunExpression(name, args, body, flags) =>
      val namedFormals = nameFormals(args)

      withEnvironmentFromDecls(namedFormals) {
        treeCopy.FunExpression(fun,
            name,
            namedFormals,
            transformExpr(body),
            flags)
      }

    case v @ Variable(name) =>
      treeCopy.Variable(v, name) withSymbol env(name)

    case EscapedVariable(v) =>
      transformExpr(v)

    case _ =>
      super.transformExpr(expression)
  }

  def extractDecls(
      declarations: List[Declaration]): (List[Variable], List[Statement]) = {
    val decls = new ListBuffer[Variable]
    val statements = new ListBuffer[Statement]

    for (declaration <- declarations) {
      declaration match {
        case variable:Variable =>
          decls += variable

        case stat @ BindStatement(left, right) =>
          decls ++= extractDeclsInExpression(left)
          statements += stat

        case stat:Statement =>
          statements += stat
      }
    }

    val namedDecls = for (v@Variable(name) <- decls.toList)
      yield treeCopy.Variable(v, name) withSymbol new VariableSymbol(name)

    (namedDecls, statements.toList)
  }

  def extractDeclsInExpression(expr: Expression) = expr match {
    case variable:Variable =>
      List(variable)

    case _ =>
      Nil<commit_msg>Better extraction of declarations in tricky nested locals.<commit_after>package org.mozartoz.bootcompiler
package transform

import scala.collection.mutable.ListBuffer

import ast._
import symtab._

object Namer extends Transformer with TransformUtils {
  type Env = Map[String, Symbol]

  private var env: Env = _

  override def apply(prog: Program) {
    withEnvironment(prog.builtins.topLevelEnvironment) {
      super.apply(prog)
    }
  }

  private def withEnvironment[A](newEnv: Env)(f: => A) = {
    val savedEnv = env
    env = newEnv
    try f
    finally env = savedEnv
  }

  private def withEnvironmentFromDecls[A](
      decls: List[Variable])(f: => A) = {
    val newEnv = (decls map (decl => decl.name -> decl.symbol))
    withEnvironment(env ++ newEnv)(f)
  }

  override def transformStat(statement: Statement) = statement match {
    case local @ LocalStatement(declarations, body) =>
      val (decls, stats) = extractDecls(declarations)
      val stat = statsAndStatToStat(stats, body)

      withEnvironmentFromDecls(decls) {
        if (decls.isEmpty) transformStat(stat)
        else treeCopy.LocalStatement(local, decls, transformStat(stat))
      }

    case _ =>
      super.transformStat(statement)
  }

  override def transformExpr(expression: Expression) = expression match {
    case local @ LocalExpression(declarations, body) =>
      val (decls, stats) = extractDecls(declarations)
      val expr = statsAndExprToExpr(stats, body)

      withEnvironmentFromDecls(decls) {
        if (decls.isEmpty) transformExpr(expr)
        else treeCopy.LocalExpression(local, decls, transformExpr(expr))
      }

    case proc @ ProcExpression(name, args, body, flags) =>
      val namedFormals = nameFormals(args)

      withEnvironmentFromDecls(namedFormals) {
        treeCopy.ProcExpression(proc,
            name,
            namedFormals,
            transformStat(body),
            flags)
      }

    case fun @ FunExpression(name, args, body, flags) =>
      val namedFormals = nameFormals(args)

      withEnvironmentFromDecls(namedFormals) {
        treeCopy.FunExpression(fun,
            name,
            namedFormals,
            transformExpr(body),
            flags)
      }

    case v @ Variable(name) =>
      treeCopy.Variable(v, name) withSymbol env(name)

    case EscapedVariable(v) =>
      transformExpr(v)

    case _ =>
      super.transformExpr(expression)
  }

  def extractDecls(
      declarations: List[Declaration]): (List[Variable], List[Statement]) = {
    val decls = new ListBuffer[Variable]
    val statements = new ListBuffer[Statement]

    def process(declarations: List[Declaration]) {
      for (declaration <- declarations) {
        declaration match {
          case variable:Variable =>
            decls += variable

          case stat @ BindStatement(left, right) =>
            decls ++= extractDeclsInExpression(left)
            statements += stat

          case stat @ LocalStatement(subDecls, subStat) =>
            val (declsInside, statsInside) = extractDecls(List(subStat))
            decls ++= declsInside
            statements += treeCopy.LocalStatement(stat, subDecls,
                statementsToStatement(statsInside))

          case stat @ CompoundStatement(subStatements) =>
            process(subStatements)

          case stat:Statement =>
            statements += stat
        }
      }
    }

    process(declarations)

    val namedDecls = for (v@Variable(name) <- decls.toList)
      yield treeCopy.Variable(v, name) withSymbol new VariableSymbol(name)

    (namedDecls, statements.toList)
  }

  def extractDeclsInExpression(expr: Expression) = expr match {
    case variable:Variable =>
      List(variable)

    case _ =>
      Nil
{'fraction_non_alphanumeric': 0.09657455382843984, 'fraction_numerical': 0.0, 'mean_word_length': 2.772529858849077, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4742762', 'n_tokens_mistral': 2020, 'n_tokens_neox': 1944, 'n_words': 456}