// fold.cm — AST optimization pass
// Constant folding, constant propagation, dead code elimination

var fold = function(ast) {
  var scopes = ast.scopes
  var nr_scopes = length(scopes)
  ast._diagnostics = []

  var type_tag_map = {
    array: "array", record: "record", text: "text",
    number: "number", blob: "blob"
  }

  var binary_ops = {
    "+": true, "-": true, "*": true, "/": true, "%": true,
    "**": true, "==": true, "!=": true, "<": true, ">": true,
    "<=": true, ">=": true, "&": true, "|": true, "^": true,
    "<<": true, ">>": true, ">>>": true, "&&": true, "||": true,
    ",": true, in: true
  }
  var unary_ops = {
    "!": true, "~": true, "-unary": true, "+unary": true, delete: true
  }
  var assign_ops = {
    assign: true, "+=": true, "-=": true, "*=": true,
    "/=": true, "%=": true, "<<=": true, ">>=": true,
    ">>>=": true, "&=": true, "^=": true, "|=": true,
    "**=": true, "&&=": true, "||=": true
  }
  var arith_ops = {
    "+": true, "-": true, "*": true, "/": true, "%": true, "**": true
  }
  var comparison_ops = {
    "==": true, "!=": true, "<": true, ">": true, "<=": true, ">=": true
  }

  // ============================================================
  // Helpers
  // ============================================================

  var is_literal = function(expr) {
    if (expr == null) return false
    var k = expr.kind
    return k == "number" || k == "text" || k == "true" || k == "false" || k == "null"
  }

  // Only intrinsics that can NEVER disrupt regardless of argument types
  var pure_intrinsics = {
    is_array: true, is_text: true, is_number: true, is_integer: true,
    is_function: true, is_logical: true, is_null: true, is_object: true,
    is_stone: true
  }

  var is_pure = function(expr) {
    if (expr == null) return true
    var k = expr.kind
    var i = 0
    var target = null
    if (k == "number" || k == "text" || k == "true" || k == "false" ||
        k == "null" || k == "name" || k == "this") return true
    if (k == "function") return true
    if (k == "!" || k == "~" || k == "-unary" || k == "+unary") {
      return is_pure(expr.expression)
    }
    if (k == "array") {
      i = 0
      while (i < length(expr.list)) {
        if (!is_pure(expr.list[i])) return false
        i = i + 1
      }
      return true
    }
    if (k == "record") {
      i = 0
      while (i < length(expr.list)) {
        if (expr.list[i].computed && !is_pure(expr.list[i].left)) return false
        if (!is_pure(expr.list[i].right)) return false
        i = i + 1
      }
      return true
    }
    if (k == "then") {
      return is_pure(expr.expression) && is_pure(expr.then) && is_pure(expr.else)
    }
    if (k == "==" || k == "!=" || k == "&&" || k == "||") {
      return is_pure(expr.left) && is_pure(expr.right)
    }
    if (k == "(") {
      target = expr.expression
      if (target != null && target.intrinsic == true && pure_intrinsics[target.name] == true) {
        i = 0
        while (i < length(expr.list)) {
          if (!is_pure(expr.list[i])) return false
          i = i + 1
        }
        return true
      }
    }
    return false
  }

  var copy_loc = function(from, to) {
    to.at = from.at
    to.from_row = from.from_row
    to.from_column = from.from_column
    to.to_row = from.to_row
    to.to_column = from.to_column
    return to
  }

  var make_number = function(val, src) {
    return copy_loc(src, {kind: "number", value: text(val), number: val})
  }

  var make_text = function(val, src) {
    return copy_loc(src, {kind: "text", value: val})
  }

  var make_bool = function(val, src) {
    if (val) return copy_loc(src, {kind: "true"})
    return copy_loc(src, {kind: "false"})
  }

  var make_null = function(src) {
    return copy_loc(src, {kind: "null"})
  }

  var is_truthy_literal = function(expr) {
    if (expr == null) return null
    var k = expr.kind
    var nv = null
    if (k == "true") return true
    if (k == "false" || k == "null") return false
    if (k == "number") {
      nv = expr.number
      if (nv == null) nv = number(expr.value)
      return nv != 0
    }
    if (k == "text") return length(expr.value) > 0
    return null
  }

  // ============================================================
  // Scope helpers
  // ============================================================

  var find_scope = function(fn_nr) {
    var i = 0
    while (i < nr_scopes) {
      if (scopes[i].function_nr == fn_nr) return scopes[i]
      i = i + 1
    }
    return null
  }

  var scope_var = function(fn_nr, name) {
    var sc = find_scope(fn_nr)
    if (sc == null) return null
    return sc[name]
  }

  var remove_scope_var = function(fn_nr, name) {
    var sc = find_scope(fn_nr)
    if (sc == null) return null
    delete sc[name]
  }

  // ============================================================
  // Pass 1: pre-scan for constants and function arities
  // ============================================================

  var const_defs = {}
  var fn_arities = {}

  var register_const = function(fn_nr, name, lit_node) {
    var key = text(fn_nr)
    if (const_defs[key] == null) const_defs[key] = {}
    const_defs[key][name] = lit_node
  }

  var get_const = function(fn_nr, name) {
    var key = text(fn_nr)
    if (const_defs[key] == null) return null
    return const_defs[key][name]
  }

  var register_arity = function(fn_nr, name, count) {
    var key = text(fn_nr)
    if (fn_arities[key] == null) fn_arities[key] = {}
    fn_arities[key][name] = count
  }

  var pre_scan_stmts = null
  var pre_scan_fn = null

  pre_scan_fn = function(node) {
    if (node == null) return null
    if (node.statements != null) pre_scan_stmts(node.statements, node.function_nr)
    if (node.disruption != null) pre_scan_stmts(node.disruption, node.function_nr)
  }

  pre_scan_stmts = function(stmts, fn_nr) {
    var i = 0
    var j = 0
    var stmt = null
    var kind = null
    var name = null
    var sv = null
    var item = null
    var rhs_target = null
    while (i < length(stmts)) {
      stmt = stmts[i]
      kind = stmt.kind
      if (kind == "def") {
        name = stmt.left.name
        if (name != null && is_literal(stmt.right)) {
          sv = scope_var(fn_nr, name)
          if (sv != null && !sv.closure) {
            register_const(fn_nr, name, stmt.right)
          }
        }
        if (name != null && stmt.right != null && stmt.right.kind == "(") {
          rhs_target = stmt.right.expression
          if (rhs_target != null && rhs_target.intrinsic == true) {
            sv = scope_var(fn_nr, name)
            if (sv != null && sv.type_tag == null) {
              if (type_tag_map[rhs_target.name] != null) sv.type_tag = type_tag_map[rhs_target.name]
            }
          }
        }
      } else if (kind == "function") {
        name = stmt.name
        if (name != null && stmt.arity != null) {
          register_arity(fn_nr, name, stmt.arity)
        }
        pre_scan_fn(stmt)
      } else if (kind == "var") {
        if (stmt.right != null && stmt.right.kind == "function" && stmt.right.arity != null) {
          name = stmt.left.name
          if (name != null) {
            sv = scope_var(fn_nr, name)
            if (sv != null && sv.make == "var") {
              register_arity(fn_nr, name, stmt.right.arity)
            }
          }
        }
      } else if (kind == "var_list") {
        j = 0
        while (j < length(stmt.list)) {
          item = stmt.list[j]
          if (item.kind == "var" && item.right != null && item.right.kind == "function" && item.right.arity != null) {
            name = item.left.name
            if (name != null) {
              sv = scope_var(fn_nr, name)
              if (sv != null && sv.make == "var") {
                register_arity(fn_nr, name, item.right.arity)
              }
            }
          }
          j = j + 1
        }
      }
      i = i + 1
    }
  }

  var pre_scan_expr_fns = null
  pre_scan_expr_fns = function(expr) {
    if (expr == null) return null
    var k = expr.kind
    var i = 0
    if (k == "function") {
      pre_scan_fn(expr)
    }
    if (expr.left != null) pre_scan_expr_fns(expr.left)
    if (expr.right != null) pre_scan_expr_fns(expr.right)
    if (expr.expression != null) pre_scan_expr_fns(expr.expression)
    if (expr.then != null) pre_scan_expr_fns(expr.then)
    if (expr.else != null) pre_scan_expr_fns(expr.else)
    if (k == "(" || k == "array") {
      i = 0
      while (i < length(expr.list)) {
        pre_scan_expr_fns(expr.list[i])
        i = i + 1
      }
    }
    if (k == "record") {
      i = 0
      while (i < length(expr.list)) {
        if (expr.list[i].computed) pre_scan_expr_fns(expr.list[i].left)
        pre_scan_expr_fns(expr.list[i].right)
        i = i + 1
      }
    }
  }

  var pre_scan_stmt_exprs = null
  pre_scan_stmt_exprs = function(stmts, fn_nr) {
    var i = 0
    var j = 0
    var stmt = null
    var kind = null
    while (i < length(stmts)) {
      stmt = stmts[i]
      kind = stmt.kind
      if (kind == "var" || kind == "def") {
        pre_scan_expr_fns(stmt.right)
      } else if (kind == "var_list") {
        j = 0
        while (j < length(stmt.list)) {
          pre_scan_expr_fns(stmt.list[j].right)
          j = j + 1
        }
      } else if (kind == "call") {
        pre_scan_expr_fns(stmt.expression)
      } else if (kind == "if") {
        pre_scan_expr_fns(stmt.expression)
        pre_scan_stmt_exprs(stmt.then, fn_nr)
        pre_scan_stmt_exprs(stmt.list, fn_nr)
        if (stmt.else != null) pre_scan_stmt_exprs(stmt.else, fn_nr)
      } else if (kind == "while" || kind == "do") {
        pre_scan_expr_fns(stmt.expression)
        pre_scan_stmt_exprs(stmt.statements, fn_nr)
      } else if (kind == "for") {
        if (stmt.init != null) {
          if (stmt.init.kind == "var" || stmt.init.kind == "def") {
            pre_scan_expr_fns(stmt.init.right)
          } else {
            pre_scan_expr_fns(stmt.init)
          }
        }
        pre_scan_expr_fns(stmt.test)
        pre_scan_expr_fns(stmt.update)
        pre_scan_stmt_exprs(stmt.statements, fn_nr)
      } else if (kind == "return" || kind == "go") {
        pre_scan_expr_fns(stmt.expression)
      } else if (kind == "block") {
        pre_scan_stmt_exprs(stmt.statements, fn_nr)
      } else if (kind == "label") {
        if (stmt.statement != null) {
          pre_scan_stmt_exprs([stmt.statement], fn_nr)
        }
      } else if (kind == "function") {
        // already handled in pre_scan_stmts
        null
      }
      i = i + 1
    }
  }

  var pre_scan = function() {
    pre_scan_stmts(ast.statements, 0)
    pre_scan_stmts(ast.functions, 0)
    pre_scan_stmt_exprs(ast.statements, 0)
    pre_scan_stmt_exprs(ast.functions, 0)
  }

  // ============================================================
  // Pass 2: fold expressions and statements
  // ============================================================

  var fold_expr = null
  var fold_stmt = null
  var fold_stmts = null
  var fold_fn = null

  fold_expr = function(expr, fn_nr) {
    if (expr == null) return null
    var k = expr.kind
    var left = null
    var right = null
    var lv = null
    var rv = null
    var result = null
    var i = 0
    var sv = null
    var lit = null
    var cond_k = null
    var ek = null
    var target = null
    var ar = null
    var akey = null
    var tv = null
    var att = null
    var arg = null

    // Recurse into children first (bottom-up)
    if (binary_ops[k] == true) {
      expr.left = fold_expr(expr.left, fn_nr)
      expr.right = fold_expr(expr.right, fn_nr)
    } else if (k == "." || k == "[") {
      expr.left = fold_expr(expr.left, fn_nr)
      if (k == "[" && expr.right != null) expr.right = fold_expr(expr.right, fn_nr)
    } else if (unary_ops[k] == true) {
      expr.expression = fold_expr(expr.expression, fn_nr)
    } else if (k == "++" || k == "--") {
      return expr
    } else if (k == "then") {
      expr.expression = fold_expr(expr.expression, fn_nr)
      expr.then = fold_expr(expr.then, fn_nr)
      expr.else = fold_expr(expr.else, fn_nr)
    } else if (k == "(") {
      expr.expression = fold_expr(expr.expression, fn_nr)
      i = 0
      while (i < length(expr.list)) {
        expr.list[i] = fold_expr(expr.list[i], fn_nr)
        i = i + 1
      }
    } else if (k == "array" || k == "text literal") {
      i = 0
      while (i < length(expr.list)) {
        expr.list[i] = fold_expr(expr.list[i], fn_nr)
        i = i + 1
      }
    } else if (k == "record") {
      i = 0
      while (i < length(expr.list)) {
        if (expr.list[i].computed) {
          expr.list[i].left = fold_expr(expr.list[i].left, fn_nr)
        }
        expr.list[i].right = fold_expr(expr.list[i].right, fn_nr)
        i = i + 1
      }
    } else if (k == "function") {
      fold_fn(expr)
      return expr
    } else if (assign_ops[k] == true) {
      expr.right = fold_expr(expr.right, fn_nr)
      return expr
    }

    // Constant propagation: name → literal
    if (k == "name" && expr.level == 0) {
      lit = get_const(fn_nr, expr.name)
      if (lit != null) {
        sv = scope_var(fn_nr, expr.name)
        if (sv != null && !sv.closure) {
          return copy_loc(expr, {kind: lit.kind, value: lit.value, number: lit.number})
        }
      }
      sv = scope_var(fn_nr, expr.name)
      if (sv != null && sv.type_tag != null) {
        expr.type_tag = sv.type_tag
      }
      return expr
    }

    // Binary constant folding
    if (arith_ops[k] == true) {
      left = expr.left
      right = expr.right
      if (left != null && right != null && left.kind == "number" && right.kind == "number") {
        lv = left.number
        rv = right.number
        if (lv == null) lv = number(left.value)
        if (rv == null) rv = number(right.value)
        if (k == "/") {
          if (rv == 0) return make_null(expr)
        }
        if (k == "%") {
          if (rv == 0) return make_null(expr)
        }
        result = null
        if (k == "+") result = lv + rv
        else if (k == "-") result = lv - rv
        else if (k == "*") result = lv * rv
        else if (k == "/") result = lv / rv
        else if (k == "%") result = lv - (trunc(lv / rv) * rv)
        else if (k == "**") result = lv ** rv
        if (result == null) return make_null(expr)
        return make_number(result, expr)
      }
      // text + text
      if (k == "+" && left != null && right != null && left.kind == "text" && right.kind == "text") {
        return make_text(left.value + right.value, expr)
      }
      return expr
    }

    // Comparison folding
    if (comparison_ops[k] == true) {
      left = expr.left
      right = expr.right
      if (left != null && right != null) {
        if (left.kind == "number" && right.kind == "number") {
          lv = left.number
          rv = right.number
          if (lv == null) lv = number(left.value)
          if (rv == null) rv = number(right.value)
          if (k == "==") return make_bool(lv == rv, expr)
          if (k == "!=") return make_bool(lv != rv, expr)
          if (k == "<") return make_bool(lv < rv, expr)
          if (k == ">") return make_bool(lv > rv, expr)
          if (k == "<=") return make_bool(lv <= rv, expr)
          if (k == ">=") return make_bool(lv >= rv, expr)
        }
        if (left.kind == "text" && right.kind == "text") {
          if (k == "==") return make_bool(left.value == right.value, expr)
          if (k == "!=") return make_bool(left.value != right.value, expr)
        }
      }
      return expr
    }

    // Bitwise folding
    if (k == "&" || k == "|" || k == "^" || k == "<<" || k == ">>") {
      left = expr.left
      right = expr.right
      if (left != null && right != null && left.kind == "number" && right.kind == "number") {
        lv = left.number
        rv = right.number
        if (lv == null) lv = number(left.value)
        if (rv == null) rv = number(right.value)
        if (k == "&") return make_number(lv & rv, expr)
        if (k == "|") return make_number(lv | rv, expr)
        if (k == "^") return make_number(lv ^ rv, expr)
        if (k == "<<") return make_number(lv << rv, expr)
        if (k == ">>") return make_number(lv >> rv, expr)
      }
      return expr
    }

    // Unary folding
    if (k == "!") {
      if (expr.expression != null) {
        ek = expr.expression.kind
        if (ek == "true") return make_bool(false, expr)
        if (ek == "false") return make_bool(true, expr)
      }
      return expr
    }
    if (k == "~") {
      if (expr.expression != null && expr.expression.kind == "number") {
        lv = expr.expression.number
        if (lv == null) lv = number(expr.expression.value)
        return make_number(~lv, expr)
      }
      return expr
    }
    if (k == "-unary") {
      if (expr.expression != null && expr.expression.kind == "number") {
        lv = expr.expression.number
        if (lv == null) lv = number(expr.expression.value)
        return make_number(0 - lv, expr)
      }
      return expr
    }

    // Ternary with literal condition
    if (k == "then") {
      tv = is_truthy_literal(expr.expression)
      if (tv == true) return expr.then
      if (tv == false) return expr.else
      return expr
    }

    // Call: stamp arity and fold intrinsic type checks
    if (k == "(") {
      target = expr.expression
      if (target != null && target.kind == "name" && target.level == 0) {
        ar = null
        akey = text(fn_nr)
        if (fn_arities[akey] != null) ar = fn_arities[akey][target.name]
        if (ar != null) expr.arity = ar
      }
      if (target != null && target.intrinsic == true && length(expr.list) == 1) {
        arg = expr.list[0]
        att = null
        if (arg.type_tag != null) {
          att = arg.type_tag
        } else if (arg.kind == "name" && arg.level == 0) {
          sv = scope_var(fn_nr, arg.name)
          if (sv != null) att = sv.type_tag
        }
        if (att != null) {
          if (target.name == "is_array") return make_bool(att == "array", expr)
          if (target.name == "is_text") return make_bool(att == "text", expr)
          if (target.name == "is_number") return make_bool(att == "number" || att == "integer", expr)
          if (target.name == "is_integer") return make_bool(att == "integer", expr)
          if (target.name == "is_function") return make_bool(att == "function", expr)
          if (target.name == "is_logical") return make_bool(att == "logical", expr)
          if (target.name == "is_null") return make_bool(att == "null", expr)
          if (target.name == "is_object") return make_bool(att == "record", expr)
          if (target.name == "length") {
            if (att == "array") expr.hint = "array_length"
            else if (att == "text") expr.hint = "text_length"
          }
        }
      }
      return expr
    }

    return expr
  }

  fold_stmt = function(stmt, fn_nr) {
    if (stmt == null) return null
    var k = stmt.kind
    var i = 0
    var sv = null
    var cond_k = null
    var ik = null
    var tv = null

    if (k == "var" || k == "def") {
      stmt.right = fold_expr(stmt.right, fn_nr)
      if (is_pure(stmt.right)) stmt.pure = true
      return stmt
    }
    if (k == "var_list") {
      i = 0
      while (i < length(stmt.list)) {
        stmt.list[i] = fold_stmt(stmt.list[i], fn_nr)
        i = i + 1
      }
      return stmt
    }
    if (k == "call") {
      stmt.expression = fold_expr(stmt.expression, fn_nr)
      return stmt
    }
    if (k == "if") {
      stmt.expression = fold_expr(stmt.expression, fn_nr)
      tv = is_truthy_literal(stmt.expression)
      if (tv == true) {
        stmt.then = fold_stmts(stmt.then, fn_nr)
        return {kind: "block", statements: stmt.then,
          at: stmt.at, from_row: stmt.from_row, from_column: stmt.from_column,
          to_row: stmt.to_row, to_column: stmt.to_column}
      }
      if (tv == false) {
        if (stmt.else != null && length(stmt.else) > 0) {
          stmt.else = fold_stmts(stmt.else, fn_nr)
          return {kind: "block", statements: stmt.else,
            at: stmt.at, from_row: stmt.from_row, from_column: stmt.from_column,
            to_row: stmt.to_row, to_column: stmt.to_column}
        }
        if (stmt.list != null && length(stmt.list) > 0) {
          return fold_stmt(stmt.list[0], fn_nr)
        }
        return null
      }
      stmt.then = fold_stmts(stmt.then, fn_nr)
      stmt.list = fold_stmts(stmt.list, fn_nr)
      if (stmt.else != null) stmt.else = fold_stmts(stmt.else, fn_nr)
      return stmt
    }
    if (k == "while") {
      stmt.expression = fold_expr(stmt.expression, fn_nr)
      if (stmt.expression.kind == "false" || stmt.expression.kind == "null") return null
      stmt.statements = fold_stmts(stmt.statements, fn_nr)
      return stmt
    }
    if (k == "do") {
      stmt.statements = fold_stmts(stmt.statements, fn_nr)
      stmt.expression = fold_expr(stmt.expression, fn_nr)
      return stmt
    }
    if (k == "for") {
      if (stmt.init != null) {
        ik = stmt.init.kind
        if (ik == "var" || ik == "def") {
          stmt.init = fold_stmt(stmt.init, fn_nr)
        } else {
          stmt.init = fold_expr(stmt.init, fn_nr)
        }
      }
      if (stmt.test != null) stmt.test = fold_expr(stmt.test, fn_nr)
      if (stmt.update != null) stmt.update = fold_expr(stmt.update, fn_nr)
      stmt.statements = fold_stmts(stmt.statements, fn_nr)
      return stmt
    }
    if (k == "return" || k == "go") {
      stmt.expression = fold_expr(stmt.expression, fn_nr)
      return stmt
    }
    if (k == "block") {
      stmt.statements = fold_stmts(stmt.statements, fn_nr)
      return stmt
    }
    if (k == "label") {
      stmt.statement = fold_stmt(stmt.statement, fn_nr)
      return stmt
    }
    if (k == "function") {
      fold_fn(stmt)
      return stmt
    }
    return stmt
  }

  fold_stmts = function(stmts, fn_nr) {
    var i = 0
    var stmt = null
    var out = []
    var sv = null
    var name = null
    while (i < length(stmts)) {
      stmt = fold_stmt(stmts[i], fn_nr)
      if (stmt == null) {
        i = i + 1
        continue
      }
      // Dead code elimination: unused pure var/def
      if (stmt.kind == "var" || stmt.kind == "def") {
        name = stmt.left.name
        if (name != null) {
          sv = scope_var(fn_nr, name)
          if (sv != null && sv.nr_uses == 0) {
            if (is_pure(stmt.right)) stmt.dead = true
            if (stmt.right != null && stmt.right.kind == "(" && stmt.right.expression != null && stmt.right.expression.name == "use") {
              push(ast._diagnostics, {
                severity: "warning",
                line: stmt.left.from_row + 1,
                col: stmt.left.from_column + 1,
                message: `unused import '${name}'`
              })
            } else if (stmt.kind == "def") {
              push(ast._diagnostics, {
                severity: "warning",
                line: stmt.left.from_row + 1,
                col: stmt.left.from_column + 1,
                message: `unused constant '${name}'`
              })
            } else {
              push(ast._diagnostics, {
                severity: "warning",
                line: stmt.left.from_row + 1,
                col: stmt.left.from_column + 1,
                message: `unused variable '${name}'`
              })
            }
          }
        }
      }
      // Dead pure call elimination: standalone pure calls with no result
      if (stmt.kind == "call" && is_pure(stmt.expression)) {
        stmt.dead = true
      }
      // Dead function elimination
      if (stmt.kind == "function" && stmt.name != null) {
        sv = scope_var(fn_nr, stmt.name)
        if (sv != null && sv.nr_uses == 0) {
          stmt.dead = true
          push(ast._diagnostics, {
            severity: "warning",
            line: stmt.from_row + 1,
            col: stmt.from_column + 1,
            message: `unused function '${stmt.name}'`
          })
        }
      }
      if (stmt.dead != true) push(out, stmt)
      i = i + 1
    }
    return out
  }

  fold_fn = function(node) {
    if (node == null) return null
    var fn_nr = node.function_nr
    if (fn_nr == null) return null
    // Fold param defaults
    var i = 0
    while (i < length(node.list)) {
      if (node.list[i].expression != null) {
        node.list[i].expression = fold_expr(node.list[i].expression, fn_nr)
      }
      i = i + 1
    }
    if (node.statements != null) node.statements = fold_stmts(node.statements, fn_nr)
    if (node.disruption != null) node.disruption = fold_stmts(node.disruption, fn_nr)
  }

  // ============================================================
  // Pass 3: cleanup scopes
  // ============================================================

  var cleanup = function() {
    var i = 0
    var sc = null
    var keys = null
    var j = 0
    var key = null
    var entry = null
    var slots = 0
    var close_slots = 0

    // Remove dead vars from scope records and recalculate slot counts
    while (i < nr_scopes) {
      sc = scopes[i]
      keys = array(sc)
      slots = 0
      close_slots = 0
      j = 0
      while (j < length(keys)) {
        key = keys[j]
        if (key != "function_nr") {
          entry = sc[key]
          if (entry != null && entry.nr_uses == 0 && entry.make != "input" && entry.make != "function") {
            delete sc[key]
          } else if (entry != null) {
            slots = slots + 1
            if (entry.closure) close_slots = close_slots + 1
          }
        }
        j = j + 1
      }
      i = i + 1
    }

    // Update nr_slots and nr_close_slots on function nodes
    var update_fn_slots = null
    update_fn_slots = function(node) {
      if (node == null) return null
      var fn_nr = node.function_nr
      if (fn_nr == null) return null
      var sc = find_scope(fn_nr)
      if (sc == null) return null
      var keys = array(sc)
      var s = 0
      var cs = 0
      var ki = 0
      var ent = null
      while (ki < length(keys)) {
        if (keys[ki] != "function_nr") {
          ent = sc[keys[ki]]
          if (ent != null) {
            s = s + 1
            if (ent.closure) cs = cs + 1
          }
        }
        ki = ki + 1
      }
      node.nr_slots = s
      node.nr_close_slots = cs
    }

    var walk_stmts_for_fns = null
    var walk_expr_for_fns = null

    walk_expr_for_fns = function(expr) {
      if (expr == null) return null
      var k = expr.kind
      var i = 0
      if (k == "function") {
        update_fn_slots(expr)
        walk_stmts_for_fns(expr.statements)
        walk_stmts_for_fns(expr.disruption)
        return null
      }
      if (expr.left != null) walk_expr_for_fns(expr.left)
      if (expr.right != null) walk_expr_for_fns(expr.right)
      if (expr.expression != null) walk_expr_for_fns(expr.expression)
      if (expr.then != null) walk_expr_for_fns(expr.then)
      if (expr.else != null) walk_expr_for_fns(expr.else)
      if (k == "(" || k == "array" || k == "text literal") {
        i = 0
        while (i < length(expr.list)) {
          walk_expr_for_fns(expr.list[i])
          i = i + 1
        }
      }
      if (k == "record") {
        i = 0
        while (i < length(expr.list)) {
          if (expr.list[i].computed) walk_expr_for_fns(expr.list[i].left)
          walk_expr_for_fns(expr.list[i].right)
          i = i + 1
        }
      }
    }

    walk_stmts_for_fns = function(stmts) {
      if (stmts == null) return null
      var i = 0
      var j = 0
      var stmt = null
      var k = null
      while (i < length(stmts)) {
        stmt = stmts[i]
        k = stmt.kind
        if (k == "function") {
          update_fn_slots(stmt)
          walk_stmts_for_fns(stmt.statements)
          walk_stmts_for_fns(stmt.disruption)
        } else if (k == "var" || k == "def") {
          walk_expr_for_fns(stmt.right)
        } else if (k == "var_list") {
          j = 0
          while (j < length(stmt.list)) {
            walk_expr_for_fns(stmt.list[j].right)
            j = j + 1
          }
        } else if (k == "call") {
          walk_expr_for_fns(stmt.expression)
        } else if (k == "if") {
          walk_expr_for_fns(stmt.expression)
          walk_stmts_for_fns(stmt.then)
          walk_stmts_for_fns(stmt.list)
          if (stmt.else != null) walk_stmts_for_fns(stmt.else)
        } else if (k == "while" || k == "do") {
          walk_expr_for_fns(stmt.expression)
          walk_stmts_for_fns(stmt.statements)
        } else if (k == "for") {
          if (stmt.init != null) {
            if (stmt.init.kind == "var" || stmt.init.kind == "def") {
              walk_expr_for_fns(stmt.init.right)
            } else {
              walk_expr_for_fns(stmt.init)
            }
          }
          walk_expr_for_fns(stmt.test)
          walk_expr_for_fns(stmt.update)
          walk_stmts_for_fns(stmt.statements)
        } else if (k == "return" || k == "go") {
          walk_expr_for_fns(stmt.expression)
        } else if (k == "block") {
          walk_stmts_for_fns(stmt.statements)
        } else if (k == "label") {
          if (stmt.statement != null) walk_stmts_for_fns([stmt.statement])
        }
        i = i + 1
      }
    }

    walk_stmts_for_fns(ast.statements)
    walk_stmts_for_fns(ast.functions)

    // Update intrinsics: collect what's still referenced
    var used_intrinsics = {}
    var collect_intrinsics = null
    var collect_expr_intrinsics = null

    collect_expr_intrinsics = function(expr) {
      if (expr == null) return null
      var k = expr.kind
      var i = 0
      if (k == "name" && expr.level == -1 && expr.name != null && expr.make != "functino") {
        used_intrinsics[expr.name] = true
      }
      if (expr.left != null) collect_expr_intrinsics(expr.left)
      if (expr.right != null) collect_expr_intrinsics(expr.right)
      if (expr.expression != null) collect_expr_intrinsics(expr.expression)
      if (expr.then != null) collect_expr_intrinsics(expr.then)
      if (expr.else != null) collect_expr_intrinsics(expr.else)
      if (k == "(" || k == "array" || k == "text literal") {
        i = 0
        while (i < length(expr.list)) {
          collect_expr_intrinsics(expr.list[i])
          i = i + 1
        }
      }
      if (k == "record") {
        i = 0
        while (i < length(expr.list)) {
          if (expr.list[i].computed) collect_expr_intrinsics(expr.list[i].left)
          collect_expr_intrinsics(expr.list[i].right)
          i = i + 1
        }
      }
      if (k == "function") {
        collect_intrinsics(expr.statements)
        collect_intrinsics(expr.disruption)
        i = 0
        while (i < length(expr.list)) {
          if (expr.list[i].expression != null) {
            collect_expr_intrinsics(expr.list[i].expression)
          }
          i = i + 1
        }
      }
    }

    collect_intrinsics = function(stmts) {
      if (stmts == null) return null
      var i = 0
      var j = 0
      var pi = 0
      var stmt = null
      var k = null
      while (i < length(stmts)) {
        stmt = stmts[i]
        k = stmt.kind
        if (k == "var" || k == "def") {
          collect_expr_intrinsics(stmt.right)
        } else if (k == "var_list") {
          j = 0
          while (j < length(stmt.list)) {
            collect_expr_intrinsics(stmt.list[j].right)
            j = j + 1
          }
        } else if (k == "call") {
          collect_expr_intrinsics(stmt.expression)
        } else if (k == "if") {
          collect_expr_intrinsics(stmt.expression)
          collect_intrinsics(stmt.then)
          collect_intrinsics(stmt.list)
          if (stmt.else != null) collect_intrinsics(stmt.else)
        } else if (k == "while" || k == "do") {
          collect_expr_intrinsics(stmt.expression)
          collect_intrinsics(stmt.statements)
        } else if (k == "for") {
          if (stmt.init != null) {
            if (stmt.init.kind == "var" || stmt.init.kind == "def") {
              collect_expr_intrinsics(stmt.init.right)
            } else {
              collect_expr_intrinsics(stmt.init)
            }
          }
          collect_expr_intrinsics(stmt.test)
          collect_expr_intrinsics(stmt.update)
          collect_intrinsics(stmt.statements)
        } else if (k == "return" || k == "go") {
          collect_expr_intrinsics(stmt.expression)
        } else if (k == "function") {
          collect_intrinsics(stmt.statements)
          collect_intrinsics(stmt.disruption)
          pi = 0
          while (pi < length(stmt.list)) {
            if (stmt.list[pi].expression != null) {
              collect_expr_intrinsics(stmt.list[pi].expression)
            }
            pi = pi + 1
          }
        } else if (k == "block") {
          collect_intrinsics(stmt.statements)
        } else if (k == "label") {
          if (stmt.statement != null) collect_intrinsics([stmt.statement])
        }
        i = i + 1
      }
    }

    collect_intrinsics(ast.statements)
    collect_intrinsics(ast.functions)

    var new_intrinsics = []
    i = 0
    while (i < length(ast.intrinsics)) {
      if (used_intrinsics[ast.intrinsics[i]] == true) {
        push(new_intrinsics, ast.intrinsics[i])
      }
      i = i + 1
    }
    ast.intrinsics = new_intrinsics
  }

  // ============================================================
  // Main
  // ============================================================

  pre_scan()

  // Pass 2: fold all statements and functions
  ast.statements = fold_stmts(ast.statements, 0)
  var fi = 0
  while (fi < length(ast.functions)) {
    fold_fn(ast.functions[fi])
    fi = fi + 1
  }

  // Remove dead top-level functions
  var live_fns = []
  var fn = null
  var fn_sv = null
  fi = 0
  while (fi < length(ast.functions)) {
    fn = ast.functions[fi]
    if (fn.name != null) {
      fn_sv = scope_var(0, fn.name)
      if (fn_sv != null && fn_sv.nr_uses == 0) {
        fn.dead = true
        push(ast._diagnostics, {
          severity: "warning",
          line: fn.from_row + 1,
          col: fn.from_column + 1,
          message: `unused function '${fn.name}'`
        })
      }
    }
    if (fn.dead != true) {
      push(live_fns, fn)
    }
    fi = fi + 1
  }
  ast.functions = live_fns

  // Pass 3: cleanup
  cleanup()

  return ast
}

return fold