kotlinlox/src/main/fr/celticinfo/lox/Resolver.kt

package fr.celticinfo.lox

/**
 * The Resolver class is used to resolve the scope of the variables.
 */
class Resolver(private val interpreter: Interpreter) : ExprVisitor<Unit>, StmtVisitor<Unit> {
     private val scopes = mutableListOf<MutableMap<String, Boolean>>()
    private var currentFunctionType = FunctionType.NONE

     init {
          scopes.add(mutableMapOf())
     }

     override fun visitBlock(stmt: Block) {
          beginScope()
          resolve(stmt.statements)
          endScope()
     }

    override fun visitClassStmt(stmt: ClassStmt) {
        declare(stmt.name)
        define(stmt.name)
    }

    override fun visitExpression(stmt: Expression) {
        resolve(stmt.expression)
    }

    override fun visitFunction(stmt: Function) {
        declare(stmt.name)
        define(stmt.name)
        resolveFunction(stmt, FunctionType.FUNCTION)
    }

    override fun visitIf(stmt: If) {
        resolve(stmt.condition)
        resolve(stmt.thenBranch)
        stmt.elseBranch?.let { resolve(it) }
    }

    override fun visitPrint(stmt: Print) {
        resolve(stmt.expression)
    }

    override fun visitReturn(stmt: Return) {
        if (currentFunctionType == FunctionType.NONE) {
            Lox.error(stmt.keyword, "Cannot return from top-level code.")
        }

        stmt.value?.let { resolve(it) }
    }

    override fun visitVar(stmt: Var) {
        declare(stmt.name)
        resolve(stmt.initializer)
        define(stmt.name)
    }

    override fun visitWhile(stmt: While) {
        resolve(stmt.condition)
        resolve(stmt.body)
    }

    override fun visitVariable(expr: Variable) {
        if (!scopes.isEmpty() && scopes.last()[expr.name.lexeme] == false) {
            Lox.error(expr.name, "Cannot read local variable in its own initializer.")
        }
        resolveLocal(expr, expr.name)
    }

    override fun visitAssign(expr: Assign) {
        resolve(expr.value)
        resolveLocal(expr, expr.name)
    }

    override fun visitBinary(expr: Binary) {
        resolve(expr.left)
        resolve(expr.right)
    }

    override fun visitCall(expr: Call) {
        resolve(expr.callee)
        expr.arguments.forEach { resolve(it) }
    }

    override fun visitGrouping(expr: Grouping) {
        resolve(expr.expression)
    }

    override fun visitLiteral(expr: Literal) {
        // Do nothing
    }

    override fun visitLogical(expr: Logical) {
        resolve(expr.left)
        resolve(expr.right)
    }

    override fun visitUnary(expr: Unary) {
        resolve(expr.right)
    }

    fun resolve(statements: List<Stmt?>) {
        for (statement in statements) {
            resolve(statement)
        }
    }

    private fun resolve(stmt: Stmt?) {
        stmt?.accept(this)
    }

    private fun resolve(expr: Expr?) {
        expr?.accept(this)
    }

    private fun beginScope() {
        scopes.add(mutableMapOf())
    }

    private fun endScope() {
        scopes.removeLast()
    }

    private fun declare(name: Token) {
        if (scopes.isEmpty()) return
        val scope = scopes.last()
        if (scope.containsKey(name.lexeme)) {
            Lox.error(name, "Variable with this name already declared in this scope.")
        }
        scope[name.lexeme] = false
    }

    private fun define(name: Token) {
        if (scopes.isEmpty()) return
        scopes.last()[name.lexeme] = true
    }

    private fun resolveLocal(expr: Expr, name: Token) {
        for (i in scopes.size - 1 downTo 0) {
            if (scopes[i].containsKey(name.lexeme)) {
                interpreter.resolve(expr, scopes.size - 1 - i)
                return
            }
        }
    }

    private fun resolveFunction(stmt: Function, type: FunctionType) {
        val enclosingFunctionType = currentFunctionType
        currentFunctionType = type

        beginScope()
        for (param in stmt.params) {
            declare(param)
            define(param)
        }
        resolve(stmt.body)
        endScope()

        currentFunctionType = enclosingFunctionType
    }
}

enum class FunctionType {
    NONE,
    FUNCTION
}