using SymbolicUtils: IRStructure, populate_ir!

include("stackdict.jl")

# NOTE:
# StackDict is an ordered container that tracks insertion order via a Vector of keys,
# allowing checkpointing by recording the last inserted key and rolling back to that
# state by popping entries in reverse order.

# NOTE:
# Codegen relies on fewer invariants than `IRStructure` in general. It is designed to work
# even when `ir.non_canonical_idxs` is non-empty. The invariants that it relies on are:
# 1. `operation(ir[idx::Integer])` is the correct operation to generate.
# 2. `symtype(ir[idx])` and `shape(ir[idx])` are the correct symtype and shape of the expression.
# 3. `outneighbors(ir.dependency_graph, idx)` corresponds in order to `arguments(ir[idx])`. In
#    case `operation(ir[idx]) isa BasicSymbolic{T}`, it is the first element in `outneighbors`,
#    followed by the arguments in order.

"""
    $TYPEDEF

Struct used for efficient codegen when coupled with a [`SymbolicUtils.IRStructure`](@ref).
Rewrites can be accessed using `.rewrites`, similar to `NameState`.

## Extended help

The following is an implementation detail and not public API.

### Fields

$TYPEDFIELDS
"""
mutable struct CodegenState{T}
    """
    The expression used as a result when codegen is finished.
    """
    const expr::Expr
    """
    An `Expr(:block)` that aliases some part of `expr`. All generated expressions are
    written to this block.
    """
    const block::Expr
    """
    The underlying `IRStructure`.
    """
    const ir::IRStructure{T}
    """
    A cache from indices of `ir` to the LHS of the assignment statement (name of the
    identifier) where it is calculated. For example, the expression `x + y` may be calculated
    as
    ```julia
    var"%1" = x
    var"%2" = y
    var"%3" = var"%1" + var"%2"
    ```
    In such a case, if we assume that the index of `x + y` in `ir` is `4`, then
    `cache[4] == Symbol("%3")`.
    """
    const cache::StackDict{Int32, Symbol}
    """
    Rewrite rules, similar to `NameState`.
    """
    const rewrites::Dict{Any, Any}
    misc_idx::Int

    function CodegenState{T}(expr, block, ir, cache, rewrites, misc_idx) where {T}
        # Verify that `block` is an `Expr(:block)`.
        @assert Meta.isexpr(block, :block)
        new{T}(expr, block, ir, cache, rewrites, misc_idx)
    end
end

"""
    $TYPEDSIGNATURES

Create a `CodegenState` where the resultant expression will be `expr`, IR will be
written to `block` and `ir` is the underlying `IRStructure`. Rewrite rules can optionally
be supplied as the last argument.
"""
function CodegenState(expr::Expr, block::Expr, ir::IRStructure{T}, rewrites = Dict()) where {T}
    CodegenState{T}(expr, block, ir, StackDict{Int32, Symbol}(), rewrites, 0)
end

"""
Key to use in the `rewrites` dictionary passed to [`fast_toexpr`](@ref) or [`CodegenState`](@ref)
for modifying the LHS of any declarations. The value should be a function which takes
`cs::CodegenState`, `idx::Integer` the integer index in `IRStructure` of the value
being generated (or `-1` if it is a miscellaneous variable), and `lhs::Symbol` the
current assignment LHS.
"""
const LHS_HOOK_KEY = :__lhs_hookₛᵧₘ

"""
    $TYPEDSIGNATURES

Generate the code in `sym` for the symbolic expression at index `idx` in `cs.ir::IRStructure`.
`sym` must be an `Expr`, `Symbol`, or similar which is valid to build an expression. It must
not be or contain other symbolic expressions.

Note that this is intended to be used to generate individual atomic expressions for symbolics.
In case a symbolic requires a larger compound expression, prefer to declare intermediates
via [`SymbolicUtils.Code.declare!`](@ref) and use them here.
"""
function codegen!(cs::CodegenState, idx::Integer, @nospecialize(sym))
    lhs = get_cse_name(length(cs.cache))
    lhs_hook = get(cs.rewrites, LHS_HOOK_KEY, nothing)
    # Declaring as `local` is necessary. Otherwise `symFunc` will
    # overwrite already-declared CSE variables.
    if lhs_hook !== nothing
        push!(cs.block.args, Expr(:local, Expr(:(=), lhs_hook(cs, idx, lhs), sym)))
    else
        push!(cs.block.args, Expr(:local, Expr(:(=), lhs, sym)))
    end
    insert!(cs.cache, idx, lhs)
    return lhs
end

"""
    $TYPEDSIGNATURES

Declare a variable named `lhs` which is assigned the value `rhs`. This is used to declare
ephemeral intermediate values in case the generated code for a symbolic is non-trivial. To
specify generated code for a symbolic, use [`SymbolicUtils.Code.codegen!`](@ref) instead.
`rhs` has the same restrictions as `sym` in that function.

Note that all-underscore identifiers can be assigned to in Julia but not used in later
expressions. Thus side-effecting expressions (such as `setindex!`) whose result is immaterial
can be generated by setting `lhs` to `:_`. The caller is responsible for ensuring the validity
of side-effecting expressions.

This function also does not check if the same `lhs` is declared twice. The variables
declared via this function are intended as temporaries, and repeated declarations will
overwrite the previous value.
"""
function declare!(cs::CodegenState, lhs, @nospecialize(rhs))
    lhs_hook = get(cs.rewrites, LHS_HOOK_KEY, nothing)
    if lhs_hook !== nothing
        push!(cs.block.args, Expr(:(=), lhs_hook(cs, -1, lhs), rhs))
    else
        push!(cs.block.args, Expr(:(=), lhs, rhs))
    end
    return lhs
end

const MISC_PREFIX = Vector{UInt8}("__miscₛᵧₘ")

"""
    $TYPEDSIGNATURES

Get a unique named miscellaneous identifier for use with [`SymbolicUtils.Code.declare!`](@ref).
"""
function get_misc_identifier(cs::CodegenState)
    name = get_numbered_symbol_with_prefix(MISC_PREFIX, cs.misc_idx)
    cs.misc_idx += 1
    return name
end

"""
    $TYPEDSIGNATURES

Get the allocator to use for a symbolic expression. If the rewrite rules contain
an entry for `ALLOCATOR_REWRITES_KEY`, will return that value and erase the rule
(since allocators cannot be reused). Otherwise, returns `default`.
"""
function get_allocator!(cs::CodegenState, @nospecialize(default))
    result = get(cs.rewrites, ALLOCATOR_REWRITES_KEY, default)
    delete!(cs.rewrites, ALLOCATOR_REWRITES_KEY)
    return result
end

const CodegenBookmarkT = Int32

"""
    $TYPEDSIGNATURES

Return a bookmark for `cs`, typically used to introduce new scopes and end them via
[`rollback!`](@ref).
"""
function bookmark(cs::CodegenState)
    return isempty(cs.cache) ? zero(Int32) : lastindex(cs.cache)
end

"""
    $TYPEDSIGNATURES

Rollback `cs` to the state identified by `bm`. This has similar semantics to `rollback!` for
[`SymbolicUtils.IRStructure`](@ref). It is typically used for scoping in `cs`.
"""
function rollback!(cs::CodegenState, bm::CodegenBookmarkT)
    last_cache_key = bm
    if iszero(last_cache_key)
        empty!(cs.cache)
    else
        k = lastindex(cs.cache)
        while k != last_cache_key
            delete!(cs.cache, k)
            k = lastindex(cs.cache)
        end
    end
    return cs
end

"""
    $TYPEDSIGNATURES

Introduce a new scope to `cs`. Some code generation paths may require inner scopes, for
example a `for` loop or `let` block. `enter_scope` can be used to introduce a new scope,
where expressions can be generated inside the block and are not accessible once the scope
is exited. This function takes the current `CodegenState`. It returns a new `CodegenState`
which contains an alias of this block and will write to it, alongside a bookmark object.
When the scope is exited, the caller is responsible for calling
[`SymbolicUtils.Code.exit_scope!`](@ref) with this returned `CodegenState` and bookmark
object.

Note that expressions generated in this scope may reference already-generated expressions
in the enclosing scope (the scope of `cs`). Calling `enter_scope` without a corresponding
`exit_scope!` is undefined behavior.
"""
function enter_scope(cs::CodegenState{T}) where {T}
    new_scope = Expr(:block)
    bm = bookmark(cs)
    scoped_cs = CodegenState{T}(new_scope, new_scope, cs.ir, cs.cache, cs.rewrites, cs.misc_idx)

    return scoped_cs, bm
end

"""
    $TYPEDSIGNATURES

End a scope introduced by [`SymbolicUtils.Code.enter_scope`](@ref). This function should
be called with the `CodegenState` and bookmark object returned by that function. It will
return an `Expr(:block)` of all code generated in this scope. Additionally, any subsequent
code generation using the `CodegenState` that `cs` was created from (the `CodegenState`
passed to `enter_scope`) will not reference expressions generated in this scope.

Note that the `Expr(:block)` returned by this function is not yet added to the outer
scope. That is the responsibility of the caller. Calling `exit_scope!` twice on the same
`cs` or with the same bookmark is undefined behavior.
"""
function exit_scope!(cs::CodegenState, bm::CodegenBookmarkT)
    result = cs.block
    rollback!(cs, bm)
    return result
end

function __default_allocator_call_expr(allocator_sym::Symbol, sz_expr::Expr)
    return Expr(:call, allocator_sym, sz_expr)
end

"""
    $TYPEDEF

Wraps an identifier already assigned to in codegen which represents a buffer that
should be written to by `with_allocator`. This pattern avoids having to do something
like

```julia
# Outer `codegen_function!` call
buffer = view(other_buffer, 1:4)
allocator = Returns(buffer)
# Inner `codegen_function!`
my_buffer = allocator((4,))
# stuff...
```

By enabling codegen to call `add_allocator!` with `ExistingBufferAllocator(:buffer)`.
This will then codegen into something like

```julia
# Outer `codegen_function!`
buffer = view(other_buffer, 1:4)
# Inner `codegen_function`
# Do stuff with `buffer`
```
"""
struct ExistingBufferAllocator
    name::Symbol
end

"""
    $TYPEDSIGNATURES

Generate an expression for calling the allocator `allocator` to store the result of `expr`.
"""
function codegen_allocator_call!(
        cs::CodegenState{T}, @nospecialize(allocator), expr::BasicSymbolic{T}, expr_idx::Integer;
        allocator_call_expr = __default_allocator_call_expr
    ) where {T}
    allocator isa ExistingBufferAllocator && return allocator.name
    allocator_sym::Symbol = if allocator isa BasicSymbolic{T}
        # If the allocator is a `Returns(buf)` expression, don't bother calling it. Just generate
        # `buf` and use it.
        @match allocator begin
            BSImpl.Term(; f) && if f === Returns || f isa DataType && f <: Returns end => begin
                allocator_idx = populate_ir!(cs.ir, allocator)
                buffer_sym = cs(cs.ir[Graphs.outneighbors(cs.ir.dependency_graph, allocator_idx)[1]])
                return codegen!(cs, expr_idx, buffer_sym)
            end
            _ => nothing
        end
        cs(allocator)
    else
        declare!(cs, get_misc_identifier(cs), allocator)
    end
    sh = shape(expr)::ShapeVecT
    sz_expr = Expr(:tuple)
    @union_split_smallvec sh for ax in sh
        push!(sz_expr.args, length(ax))
    end
    return codegen!(cs, expr_idx, allocator_call_expr(allocator_sym, sz_expr))
end

"""
    $TYPEDSIGNATURES

Replacement of [`SymbolicUtils.Code.toexpr`](@ref) using IR-based codegen.
"""
function fast_toexpr(sym::BasicSymbolic{T}, rewrites::Dict{Any, Any}) where {T}
    ir = IRStructure{T}()
    return fast_toexpr(sym, ir, rewrites)
end

function fast_toexpr(sym, ir::IRStructure{T}, rewrites::Dict{Any, Any}) where {T}
    expr = block = Expr(:block)
    state = CodegenState(expr, block, ir, rewrites)
    state(sym)
    return expr
end

"""
    $TYPEDSIGNATURES
"""
function fast_toexpr(sym::CodegenPrimitive, ir::IRStructure{T}, rewrites::Dict{Any, Any}) where {T}
    expr = block = Expr(:block)
    state = CodegenState(expr, block, ir, rewrites)
    lhs = state(sym)::Symbol
    for line in expr.args
        if Meta.isexpr(line, :(=)) && line.args[1] === lhs
            return line.args[2]
        end
    end
    _unreachable()
end

"""
    function codegen_function!(fn::F, cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {F, T}

Generate code for `expr` whose operation is `fn` and index in `cs.ir` is `expr_idx`.
Implementers of this function should _only_ dispatch on `F`, and the rest of the types should
be as provided in the signature above.

This function should generate a Julia expression `result` for `expr`, and call
[`SymbolicUtils.Code.codegen!`](@ref)`(cs, expr_idx, result)`. Temporary intermediate values
may be declared via [`SymbolicUtils.Code.declare!`](@ref). This function should return the
value returned by `SymbolicUtils.Code.codegen!`. Any other return value is undefined behavior.

In case the code generation process here requires the results of other symbolic expressions,
it should call `cs` to obtain them. For example, if it requires the generated code for
`arguments(expr)[1]`, it should call `cs(arguments(expr)[1])` to obtain it. This will return
a `Symbol` indicating the identifier in the generated code which is assigned the requested
value.

This function is allowed to non-destructively mutate `cs.ir`. All information present in
`cs.ir` upon entry to this function should be present when it returns. The function may
add expressions to `cs.ir` and choose to use [`SymbolicUtils.rollback!`](@ref) to restore
to an earlier state, but this earlier state should not erase information that was not
added by this function.

In case the code generation process requires creating a new scope such as a `for` loop or
`let` block, it should use [`SymbolicUtils.Code.enter_scope`](@ref) and
[`SymbolicUtils.Code.exit_scope!`](@ref).

The code generation in this function does not need to handle checking if `expr` is in
`cs.rewrites`. That functionality is handled elsewhere.
"""
function codegen_function! end

function codegen_function!(::Type{ArrayOp{T}}, cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    _allocator = get_allocator!(cs, zeros)
    output_buffer = codegen_allocator_call!(cs, _allocator, expr, expr_idx)

    # Use graph instead of destructuring `expr`
    nbors = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    innerexpr = SymbolicUtils.get_canonical_expr!(cs.ir, nbors[2])::BasicSymbolic{T}
    reducer = unwrap_const(cs.ir[nbors[3]])
    ranges = unwrap_const(cs.ir[nbors[5]])::SymbolicUtils.RangesT{T}

    new_ranges = RangesT{T}()
    new_expr = unidealize_indices(innerexpr, ranges, new_ranges)

    scoped_cs, bm = enter_scope(cs)

    inner_result = scoped_cs(new_expr)

    loopvar_order = BasicSymbolic{T}[]
    index_expr = Expr(:call, CartesianIndex)

    # `output_idxs` can be empty/all integers/symbolic. The first two cases are captured by
    # the `Const`, where all loops are reduction loops and we output a scalar or singleton
    # array. If we have parallel loops, then `output_idxs` will contain symbolic entries and
    # it will thus be present as an `array_literal`.
    out_idx_sym = cs.ir[nbors[1]]
    @match out_idx_sym begin
        BSImpl.Const(; val) && if val isa SymbolicUtils.OutIdxT{T} end => begin
            # If this is a `Const`, then none of the elements in `output_idx` are symbolic,
            # so we never populate `loopvar_order`.
            output_idx = val::SymbolicUtils.OutIdxT{T}
            @union_split_smallvec output_idx for i in output_idx
                push!(index_expr.args, scoped_cs(i))
                if i isa BasicSymbolic{T}
                    push!(loopvar_order, i)
                end
            end
        end
        BSImpl.Term(; f) && if f === SymbolicUtils.array_literal end => begin
            # Get the neighbors of the `array_literal`.
            out_idx_nbors = Graphs.outneighbors(cs.ir.dependency_graph, nbors[1])
            # The first neighbor is the size tuple, so ignore it.
            for out_idx_i in Iterators.drop(out_idx_nbors, 1)
                out_idx = cs.ir[out_idx_i]
                @match out_idx begin
                    # Constant index - singleton dimension
                    BSImpl.Const(; val) => push!(index_expr.args, val::Int)
                    # Symbolic index - parallel loop
                    _ => begin
                        push!(index_expr.args, scoped_cs(out_idx))
                        push!(loopvar_order, out_idx)
                    end
                end
            end
        end
    end
    reverse!(loopvar_order)
    ref_expr = Expr(:ref, output_buffer, index_expr)

    declare!(scoped_cs, :__accum, Expr(:call, reducer, ref_expr, inner_result))
    declare!(scoped_cs, :_, Expr(:(=), ref_expr, :__accum))

    merge!(new_ranges, ranges)
    append!(loopvar_order, keys(new_ranges))
    unique!(loopvar_order)

    outer_expr = Expr(:block)
    cur_expr = outer_expr
    for lv in loopvar_order
        new_cur_expr = Expr(:block)
        push!(cur_expr.args, Expr(:for, Expr(:(=), scoped_cs(lv), new_ranges[lv]), new_cur_expr))
        cur_expr = new_cur_expr
    end

    scoped_block = exit_scope!(scoped_cs, bm)
    push!(cur_expr.args, scoped_block)

    declare!(cs, :_, outer_expr)

    return output_buffer
end

function codegen_function!(f::SymbolicUtils.Fill, cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    _allocator = get_allocator!(cs, nothing)
    result = Expr(:call)
    if _allocator === nothing
        push!(result.args, fill)
    else
        push!(result.args, fill!)
        output_buffer = codegen_allocator_call!(cs, _allocator, expr, expr_idx)
        push!(result.args, output_buffer)
    end

    push!(result.args, cs(cs.ir[Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)[1]]))
    if _allocator === nothing
        for ax in f.sh
            push!(result.args, length(ax))
        end
        return codegen!(cs, expr_idx, result)
    end
    declare!(cs, get_misc_identifier(cs), result)
    return output_buffer
end

function write_constant_array!(cs::CodegenState{T}, vw::Symbol, val::AbstractArray) where {T}
    for (i, v) in enumerate(val)
        write_target = Expr(:ref, vw, i)
        declare!(cs, write_target, v)
    end
end

@generated function fill_arr!(buffer::AbstractArray, ::Val{sz}, args...) where {sz}
    @assert prod(sz) == length(args)
    result = Expr(:block)
    for (i, idx) in enumerate(CartesianIndices(sz))
        push!(result.args, :(buffer[$(Tuple(idx)...)] = args[$i]))
    end
    push!(result.args, :(return buffer))
    return result
end

@generated function fill_sarr(::Val{sz}, args...) where {sz}
    @assert prod(sz) == length(args)
    if length(sz) == 1
        return :(SVector{$(sz[1])}(args...))
    elseif length(sz) == 2
        return :(SMatrix{$(sz[1]), $(sz[2])}(args...))
    else
        error("fill_sarr only supports 1D and 2D arrays, got shape $sz")
    end
end

"""
Maximum total element count for which `fill_sarr` (returning a StaticArray)
will be used instead of the mutable `fill_arr!`.
"""
const STATIC_ARRAY_LIMIT::Int = 4

"""
Limit until which the generated code for an `ArrayMaker`/`array_literal` will use `fill_arr!`
over the standard codegen.
"""
const FILL_ARR_LIMIT = 16

@noinline function get_allocator_and_fill_zero!(allocator, sz)
    buffer = allocator(sz)
    fill!(buffer, 0)
    return buffer
end

function __allocator_call_and_fill_expr(allocator_sym::Symbol, sz_expr::Expr)
    return Expr(:call, get_allocator_and_fill_zero!, allocator_sym, sz_expr)
end

function __is_fill_zero(ex::BasicSymbolic{T}) where {T}
    @match ex begin
        BSImpl.ArrayOp(; term) && if term !== nothing end => @match term begin
            BSImpl.Term(; f, args) => f isa SymbolicUtils.Fill && SymbolicUtils._iszero(args[1])
            _ => false
        end
        _ => false
    end
end

function __allocator_is_returns_expr(::Type{T}, @nospecialize(ex)) where {T}
    if ex isa BasicSymbolic{T}
        @match ex begin
            BSImpl.Term(; f) => return f === Returns || f isa DataType && f <: Returns
            _ => return false
        end
    end
    return false
end

@noinline function batched_setindex!(buffer::AbstractArray{T, D}, values::NTuple{N, T}, idxs::NTuple{N, CartesianIndex{D}}) where {T, D, N}
    @inbounds for i in 1:N
        buffer[idxs[i]] = values[i]
    end
end
@noinline function batched_setindex!(buffer::AbstractArray{T, D}, value::T, idxs::AbstractArray{CartesianIndex{D}}) where {T, D}
    @inbounds for idx in idxs
        buffer[idx] = value
    end
end

BATCHED_SETINDEX_BATCH_SIZE::Int = 64

function codegen_function!(::Type{ArrayMaker{T}}, cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    _allocator = get_allocator!(cs, zeros)
    len = length(expr)::Int

    # We can only rely on the graph during codegen, so this is a roundabout way of obtaining
    # the `regions` and `values` fields of the `ArrayMaker`.
    expr_nbors = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    regions = unwrap_const(cs.ir[expr_nbors[1]])::SymbolicUtils.RegionsT
    values_expr_idx = expr_nbors[2]
    values_args = Graphs.outneighbors(cs.ir.dependency_graph, values_expr_idx)
    # Indices of the expressions in `values`.
    values_exprs_idxs = collect(Iterators.drop(values_args, 1))
    sh = shape(expr)

    # The array is small enough, so scalarize it and use `fill_arr!` or `fill_sarr`.
    if len <= FILL_ARR_LIMIT
        expr = SymbolicUtils.get_canonical_expr!(cs.ir, expr_idx)
        sz_expr = Expr(:tuple)
        for ax in sh
            push!(sz_expr.args, length(ax))
        end
        if all(<=(STATIC_ARRAY_LIMIT), sz_expr.args) && length(sz_expr.args) <= 2 && _allocator === zeros
            result = Expr(:call, fill_sarr, Expr(:call, Val, sz_expr))
            for idx in SymbolicUtils.stable_eachindex(expr)
                push!(result.args, cs(expr[idx]))
            end
        else
            output_buffer = codegen_allocator_call!(cs, _allocator, expr, expr_idx)
            result = Expr(:call, fill_arr!, output_buffer)
            push!(result.args, Expr(:call, Val, sz_expr))
            for idx in SymbolicUtils.stable_eachindex(expr)
                push!(result.args, cs(expr[idx]))
            end
        end
        return declare!(cs, get_misc_identifier(cs), result)
    end

    if _allocator !== zeros && !__allocator_is_returns_expr(T, _allocator) &&
            isequal(regions[1], sh) && __is_fill_zero(cs.ir[first(values_exprs_idxs)])
        output_buffer = codegen_allocator_call!(
            cs, _allocator, expr, expr_idx;
            allocator_call_expr = __allocator_call_and_fill_expr
        )
        ndrop = 1
    else
        output_buffer = codegen_allocator_call!(cs, _allocator, expr, expr_idx)
        ndrop = 0
    end

    # Indices and values of regions that write to a scalar
    scalar_idxs = Dict{ShapeVecT, Any}()
    D = length(sh)::Int
    # Utility buffers
    one_regions = ShapeVecT[]
    negone_regions = ShapeVecT[]
    other_regions = ShapeVecT[]
    other_values = Any[]

    eltype_expr = declare!(cs, get_misc_identifier(cs), Expr(:call, eltype, output_buffer))
    one_expr = Expr(:call, one, eltype_expr)
    negone_expr = Expr(:call, :(-), one_expr)

    function materialize_scalar_writes()
        empty!(one_regions)
        empty!(negone_regions)
        empty!(other_regions)
        empty!(other_values)
        # Batch all writes of +-1 together
        for region in keys(scalar_idxs)
            _val = scalar_idxs[region]
            if _val === 1
                push!(one_regions, region)
            elseif _val === -1
                push!(negone_regions, region)
            else
                push!(other_regions, region)
                push!(other_values, _val)
            end
        end
        if !isempty(one_regions)
            idxs_buffer = Vector{CartesianIndex{D}}(map(CartesianIndex{D} ∘ Tuple ∘ Base.Fix1(map, first), one_regions))
            # Directly interpolating an array into an expression avoids the allocation.
            declare!(cs, :_, Expr(:call, batched_setindex!, output_buffer, one_expr, idxs_buffer))
        end
        if !isempty(negone_regions)
            idxs_buffer = Vector{CartesianIndex{D}}(map(CartesianIndex{D} ∘ Tuple ∘ Base.Fix1(map, first), negone_regions))
            declare!(cs, :_, Expr(:call, batched_setindex!, output_buffer, negone_expr, idxs_buffer))
        end
        # Use `Iterators.partition` to avoid creating overly large tuples
        for (regs, vals) in zip(
                Iterators.partition(other_regions, BATCHED_SETINDEX_BATCH_SIZE),
                Iterators.partition(other_values, BATCHED_SETINDEX_BATCH_SIZE)
            )
            vals_expr = Expr(:tuple)
            for val in vals
                push!(vals_expr.args, Expr(:call, eltype_expr, val))
            end
            idxs_expr = Expr(:tuple)
            for reg in regs
                push!(idxs_expr.args, CartesianIndex{D}(Tuple(map(first, reg))))
            end
            declare!(cs, :_, Expr(:call, batched_setindex!, output_buffer, vals_expr, idxs_expr))
        end
    end

    for cur_region in Iterators.drop(eachindex(regions), ndrop)
        region = regions[cur_region]
        val_idx = values_exprs_idxs[cur_region]
        val = cs.ir[val_idx]
        # We are writing to a single element
        if all(isone ∘ length, region)
            # The goal is to batch scalar writes, so we defer them. Storing them in a
            # `Dict` allows later writes to overwrite earlier ones without additional
            # codegen.
            @match val begin
                BSImpl.Const(; val = _val) => begin
                    _inner = first(_val)
                    if _isone(_inner)
                        scalar_idxs[region] = 1
                    elseif _isone(-_inner)
                        scalar_idxs[region] = -1
                    else
                        scalar_idxs[region] = _inner
                    end
                end
                BSImpl.Term(; f) && if f === SymbolicUtils.array_literal end => begin
                    first_val_idx = Graphs.outneighbors(cs.ir.dependency_graph, val_idx)[2]
                    scalar_idxs[region] = cs(cs.ir[first_val_idx])
                end
                BSImpl.Term(; f) && if f isa Union{DataType, UnionAll} && f <: StaticArraysCore.StaticArray end => begin
                    first_val_idx = Graphs.outneighbors(cs.ir.dependency_graph, val_idx)[1]
                    scalar_idxs[region] = cs(cs.ir[first_val_idx])
                end
                _ => begin
                    # Unfortunately, scalarizing `val` requires getting the canonical expr.
                    canon_ex = SymbolicUtils.get_canonical_expr!(cs.ir, val_idx)
                    scalar_idxs[region] = cs(canon_ex[SymbolicUtils.stable_eachindex(canon_ex)[1]])
                end
            end
            continue
        end

        # We are writing to a block. First, materialize all scalar writes that happened until now.
        # This is operating on the assumption that non-scalar writes conflict with all scalar writes.
        # If this case is hit frequently, it might be useful to check for actual conflicts.
        if !isempty(scalar_idxs)
            materialize_scalar_writes()
        end
        empty!(scalar_idxs)
        # Materialize the view into the output buffer that we will write to.
        view_expr = Expr(:call, view, output_buffer)
        @union_split_smallvec region append!(view_expr.args, region)
        vw = declare!(cs, get_misc_identifier(cs), view_expr)
        @match val begin
            # Replace the allocator. We want it to write directly to the output
            # No point constructing a new expression, just do what `with_allocator`
            # would do in `codegen_ir!`.
            BSImpl.Term(; f) && if f === with_allocator end => begin
                old_allocator = add_allocator!(cs, ExistingBufferAllocator(vw))
                val_nbors = Graphs.outneighbors(cs.ir.dependency_graph, val_idx)
                cs(cs.ir[val_nbors[2]])
                reset_allocator!(cs, old_allocator)
            end
            # Any other operation that supports `with_allocator` should be generated to
            # write to the buffer directly.
            _ && if supports_with_allocator(val) end => begin
                old_allocator = add_allocator!(cs, ExistingBufferAllocator(vw))
                cs(val)
                reset_allocator!(cs, old_allocator)
            end
            # Constant arrays can special-case
            BSImpl.Const(; val) => write_constant_array!(cs, vw, val)
            # We don't have a fallback, so materialize the array corresponding to `val`
            # and `copyto!` it into the output buffer.
            _ => begin
                temp_array = cs(val)
                declare!(cs, get_misc_identifier(cs), Expr(:call, copyto!, vw, temp_array))
            end
        end
    end
    if !isempty(scalar_idxs)
        materialize_scalar_writes()
    end
    return declare!(cs, get_misc_identifier(cs), output_buffer)
end

function codegen_function!(
        ::typeof(SymbolicUtils.array_literal), cs::CodegenState{T}, expr::BasicSymbolic{T},
        expr_idx::Integer
    ) where {T}
    _allocator = get_allocator!(cs, nothing)
    if _allocator === nothing
        return codegen_function_no_nanmath!(SymbolicUtils.array_literal, cs, expr, expr_idx)
    end
    args_idxs = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    sh = @match expr begin
        BSImpl.Term(; shape) => shape
    end
    len = length(expr)::Int

    if len <= FILL_ARR_LIMIT
        sz_val = unwrap_const(cs.ir[first(args_idxs)])
        ndims_arr = length(sz_val)
        if all(<=(STATIC_ARRAY_LIMIT), sz_val) && ndims_arr <= 2 && _allocator === zeros
            result = Expr(:call, fill_sarr, Expr(:call, Val, sz_val))
            for arg_idx in Iterators.drop(args_idxs, 1)
                push!(result.args, cs(cs.ir[arg_idx]))
            end
        else
            output_buffer = codegen_allocator_call!(cs, _allocator, expr, expr_idx)
            result = Expr(:call, fill_arr!, output_buffer, Expr(:call, Val, sz_val))
            for arg_idx in Iterators.drop(args_idxs, 1)
                push!(result.args, cs(cs.ir[arg_idx]))
            end
        end
        return declare!(cs, get_misc_identifier(cs), result)
    end
    output_buffer = codegen_allocator_call!(cs, _allocator, expr, expr_idx)

    eltype_expr = declare!(cs, get_misc_identifier(cs), Expr(:call, eltype, output_buffer))
    cart_idxs = CartesianIndices(Tuple(sh::ShapeVecT))
    for iter in Iterators.partition(enumerate(Iterators.drop(args_idxs, 1)), BATCHED_SETINDEX_BATCH_SIZE)
        vals_expr = Expr(:tuple)
        idxs_expr = Expr(:tuple)
        for (idx, arg_idx) in iter
            push!(vals_expr.args, Expr(:call, eltype_expr, cs(cs.ir[arg_idx])))
            push!(idxs_expr.args, cart_idxs[idx])
        end
        declare!(cs, :_, Expr(:call, batched_setindex!, output_buffer, vals_expr, idxs_expr))
    end

    return output_buffer
end

"""
Somewhat arbitrary limit for the maximum number of arguments when calling n-ary
functions such as `+` or `*`.
"""
const NARY_CALL_LIMIT = 12

function codegen_function!(@nospecialize(op::Union{typeof(*), typeof(+)}), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    if get(cs.rewrites, :sort_addmul, true)::Bool
        @assert isempty(cs.ir.non_canonical_idxs) "Sorted add/mul requires canonical IRStructure"
        args_idxs = Int32[]
        sargs = sorted_arguments(expr)
        for arg in sargs
            push!(args_idxs, cs.ir[arg])
        end
        # For type-stability, so it matches with the other branch
        args_idxs = SymbolicUtils.AdjView{Int32}(args_idxs)
    else
        args_idxs = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    end

    if length(args_idxs) == 1
        return codegen!(cs, expr_idx, Expr(:call, op, cs(cs.ir[first(args_idxs)])))
    elseif length(args_idxs) == 2
        return codegen!(cs, expr_idx, Expr(:call, op, cs(cs.ir[args_idxs[1]]), cs(cs.ir[args_idxs[2]])))
    end

    if symtype(expr) <: Number
        cur = Expr(:call, op)
        for (i, idx) in enumerate(args_idxs)
            if i % NARY_CALL_LIMIT == 0
                cur = Expr(:call, op, cur, cs(cs.ir[idx]))
            else
                push!(cur.args, cs(cs.ir[idx]))
            end
        end
        return codegen!(cs, expr_idx, cur)
    end

    # This is tensor add/mul. Do a `foldl`-esque implementation reinserting into the IR. This
    # CSEs identical matrix multiplications reliably. `A * B * C` and `A * B * D` will
    # reuse the common prefix result. Reuse in addition relies on hash ordering.
    prefix_args = ArgsT{T}()
    resize!(prefix_args, length(args_idxs) - 1)
    for i in 1:length(prefix_args)
        prefix_args[i] = cs.ir[args_idxs[i]]
    end
    # Yes, we could call `promote_symtype` and `promote_shape` but those don't really make
    # a difference in codegen and the entire point is to reuse common prefixes. Reusing
    # `symtype` and `shape` of `expr` is problematic because of `Matrix * Vector` et. al.
    # All we want is to recurse back into this method and for identical prefixes to hashcons
    # identically, which this accomplishes.
    prefix_term = BSImpl.Term{T}(op, prefix_args; type = Any, shape = Unknown(-1))
    prefix_expr = cs(prefix_term)
    last_term = cs.ir[args_idxs[length(args_idxs)]]
    last_expr = cs(last_term)
    return codegen!(cs, expr_idx, Expr(:call, op, prefix_expr, last_expr))
end

function codegen_function!(op::typeof(^), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    base_idx, exp_idx = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    base_sym = cs(cs.ir[base_idx])
    exp = cs.ir[exp_idx]
    @match exp begin
        BSImpl.Const(; val) => begin
            if val isa Real
                if _isone(val)
                    return codegen!(cs, expr_idx, base_sym)
                end
                if _isone(-val)
                    return codegen!(cs, expr_idx, Expr(:call, inv, base_sym))
                end
                if !(val isa Integer) && get(cs.rewrites, :nanmath, false)::Bool
                    op = NaNMath.pow
                end
            end
            return codegen!(cs, expr_idx, Expr(:call, op, base_sym, val))
        end
        _ => begin
            sT = symtype(exp)
            exp_sym = cs(exp)
            if get(cs.rewrites, :nanmath, false)::Bool
                return codegen!(cs, expr_idx, Expr(:call, NaNMath.pow, base_sym, exp_sym))
            else
                return codegen!(cs, expr_idx, Expr(:call, ^, base_sym, exp_sym))
            end
        end
    end
end

function codegen_function!(::typeof(ifelse), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    cond_idx, true_idx, false_idx = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    return codegen!(cs, expr_idx, Expr(:if, cs(cs.ir[cond_idx]), cs(cs.ir[true_idx]), cs(cs.ir[false_idx])))
end

# `ifelse_branching` emits each branch body into its own codegen scope so it lands inside the
# corresponding `if`/`else` arm rather than being hoisted ahead of the conditional. Together
# with `cse_inside_expr(_, ::typeof(ifelse_branching)) = false`, this guarantees the untaken
# branch is never evaluated, even when CSE is enabled. (`ifelse_eager` needs no method here:
# it lowers via the generic `codegen_function!` fallback to an eager call.)
function codegen_function!(::typeof(ifelse_branching), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    cond_idx, true_idx, false_idx = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    cond = cs(cs.ir[cond_idx])

    true_cs, true_bm = enter_scope(cs)
    true_val = true_cs(cs.ir[true_idx])
    true_block = exit_scope!(true_cs, true_bm)
    push!(true_block.args, true_val)

    false_cs, false_bm = enter_scope(cs)
    false_val = false_cs(cs.ir[false_idx])
    false_block = exit_scope!(false_cs, false_bm)
    push!(false_block.args, false_val)

    return codegen!(cs, expr_idx, Expr(:if, cond, true_block, false_block))
end

function codegen_function_no_nanmath!(@nospecialize(op), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    result = Expr(:call)
    if !(op isa BasicSymbolic{T})
        push!(result.args, op)
    end
    # If the operation is symbolic, we'll push it as the function to call here.
    # If it isn't symbolic, we'll push it above and here we add the arguments.
    nbors = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    for nbor in nbors
        push!(result.args, cs(cs.ir[nbor]))
    end
    return codegen!(cs, expr_idx, result)
end

function codegen_function!(@nospecialize(op), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    # Avoid using `op in NaNMathFuns` since that is dynamic dispatch.
    # Also avoid using `nameof(op)` for the same reason.
    if !get(cs.rewrites, :nanmath, false)::Bool
        return codegen_function_no_nanmath!(op, cs, expr, expr_idx)
    elseif op === sin
        name = :sin
    elseif op === cos
        name = :cos
    elseif op === tan
        name = :tan
    elseif op === asin
        name = :asin
    elseif op === acos
        name = :acos
    elseif op === acosh
        name = :acosh
    elseif op === atanh
        name = :atanh
    elseif op === log
        name = :log
    elseif op === log2
        name = :log2
    elseif op === log10
        name = :log10
    elseif op === lgamma
        name = :lgamma
    elseif op === log1p
        name = :log1p
    elseif op === sqrt
        name = :sqrt
    else
        return codegen_function_no_nanmath!(op, cs, expr, expr_idx)
    end
    op = GlobalRef(NaNMath, name)
    return codegen_function_no_nanmath!(op, cs, expr, expr_idx)
end

function codegen_function!(::typeof(getindex), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    args_idxs = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    if isequal(cs.ir[first(args_idxs)], SymbolicUtils.idxs_for_arrayop(T))
        offset = unwrap_const(cs.ir[args_idxs[2]])::Int
        return codegen_arrayop_indexer(offset)
    end
    result = Expr(:ref)
    for idx in args_idxs
        push!(result.args, cs(cs.ir[idx]))
    end
    return codegen!(cs, expr_idx, result)
end

function codegen_function!(@nospecialize(op::SymbolicUtils.Mapper), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    result = Expr(:call, map)
    fn = op.f
    if fn isa BasicSymbolic{T}
        push!(result.args, cs(fn))
    else
        push!(result.args, fn)
    end
    args_idxs = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    for idx in args_idxs
        push!(result.args, cs(cs.ir[idx]))
    end
    return codegen!(cs, expr_idx, result)
end

function codegen_function!(@nospecialize(op::SymbolicUtils.Mapreducer), cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    result = Expr(:call, mapreduce)
    kws = Expr(:parameters)
    if op.dims isa Int
        dims = op.dims
        if dims isa BasicSymbolic{T}
            push!(kws.args, Expr(:kw, :dims, cs(dims)))
        else
            push!(kws.args, Expr(:kw, :dims, dims))
        end
    end
    if op.init !== nothing
        init = op.init
        if init isa BasicSymbolic{T}
            push!(kws.args, Expr(:kw, :init, cs(init)))
        else
            push!(kws.args, Expr(:kw, :init, init))
        end
    end
    if !isempty(kws.args)
        push!(result.args, kws)
    end
    fn = op.f
    if fn isa BasicSymbolic{T}
        push!(result.args, cs(fn))
    else
        push!(result.args, fn)
    end
    reducer = op.reduce
    if reducer isa BasicSymbolic{T}
        push!(result.args, cs(reducer))
    else
        push!(result.args, reducer)
    end
    args_idxs = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    for idx in args_idxs
        push!(result.args, cs(cs.ir[idx]))
    end
    return codegen!(cs, expr_idx, result)
end

function add_allocator!(cs::CodegenState, @nospecialize(alloc))
    old_allocator = get(cs.rewrites, ALLOCATOR_REWRITES_KEY, nothing)
    cs.rewrites[ALLOCATOR_REWRITES_KEY] = alloc
    return old_allocator
end

function reset_allocator!(cs::CodegenState, @nospecialize(old_alloc))
    cs.rewrites[ALLOCATOR_REWRITES_KEY] = old_alloc
    return nothing
end

function codegen_ir!(cs::CodegenState{T}, idx::Integer) where {T}
    cached = get(cs.cache, idx, nothing)
    if cached isa Symbol
        return cached
    end
    ir = cs.ir
    sym = ir[idx]
    if !SymbolicUtils.isconst(sym)
        new_expr = get(cs.rewrites, sym, nothing)
        if new_expr isa BasicSymbolic{T}
            new_idx = populate_ir!(ir, new_expr)
            return codegen_ir!(cs, new_idx)
        elseif new_expr isa Expr
            return codegen!(cs, idx, new_expr)
        elseif new_expr isa Symbol
            return codegen!(cs, idx, new_expr)
        elseif new_expr !== nothing
            _unreachable()
        end
    end

    @match sym begin
        BSImpl.Const(; val) => if val isa CodegenPrimitive
            return cs(val)
        else
            return codegen!(cs, idx, val)
        end
        BSImpl.Sym(; name) => return codegen!(cs, idx, name)
        BSImpl.ArrayOp(; term) => if term isa BasicSymbolic{T}
            # NOTE: This used to call `populate_ir!(ir, term)` and
            # then `codegen_ir!` on that. However, if the IR is non-canonical
            # this leads to incorrect results since `term` is incorrect
            # and we should use the graph at this node (which reflects what
            # `term` should be). `codegen_function!` overcomes this by doing
            # the "normal" thing with the correct operation.
            return codegen_function!(operation(term), cs, sym, idx)
        else
            return codegen_function!(ArrayOp{T}, cs, sym, idx)
        end
        BSImpl.ArrayMaker(;) => return codegen_function!(ArrayMaker{T}, cs, sym, idx)
        BSImpl.Term(; f) => begin
            if f === (^)
                return codegen_function!(^, cs, sym, idx)
            elseif f === ifelse
                return codegen_function!(ifelse, cs, sym, idx)
            elseif f === getindex
                return codegen_function!(getindex, cs, sym, idx)
            elseif f === with_allocator
                args_idxs = Graphs.outneighbors(ir.dependency_graph, idx)
                new_allocator = ir[first(args_idxs)]
                old_allocator = add_allocator!(cs, new_allocator)
                result = cs(ir[args_idxs[2]])
                reset_allocator!(cs, old_allocator)
                return result
            else
                return codegen_function!(f, cs, sym, idx)::Symbol
            end
        end
        BSImpl.Div(;) => begin
            num_idx, den_idx = Graphs.outneighbors(ir.dependency_graph, idx)
            return codegen!(cs, idx, Expr(:call, /, cs(ir[num_idx]), cs(ir[den_idx])))
        end
        BSImpl.AddMul(; variant) => if variant === AddMulVariant.ADD
            return codegen_function!((+), cs, sym, idx)
        else
            return codegen_function!((*), cs, sym, idx)
        end
    end
end

function (cs::CodegenState)(arr::AbstractArray)
    uarr = unwrap(arr)
    if uarr !== arr
        return cs(uarr)
    end
    return cs(MakeArray(arr, typeof(arr)))
end

function (cs::CodegenState)(arr::Union{SparseVector, SparseMatrixCSC})
    return cs(MakeSparseArray(arr))
end

function (cs::CodegenState)(@nospecialize(thing))
    uthing = unwrap(thing)
    if uthing !== thing
        return cs(uthing)
    end
    return declare!(cs, get_misc_identifier(cs), thing)
end

function (cs::CodegenState)(expr::BasicSymbolic{T})::Symbol where {T}
    idx = populate_ir!(cs.ir, expr)
    codegen_ir!(cs, idx)
end

function (cs::CodegenState)(expr::Assignment)
    lhs = manual_dispatch_toexpr(expr.lhs, NameState(cs.rewrites))
    # Previously this only generated a block if `Meta.isexpr(lhs, :call)` for inline
    # function declarations. This was changed when CSE variables were made to declare
    # as `local`, since `let` block's don't allow `local` for assignments in the header.
    scs, bm = enter_scope(cs)
    scs(expr.rhs)
    rhs = exit_scope!(scs, bm)
    
    return declare!(cs, lhs, rhs)
end

const ARGNAME_PREFIX = Vector{UInt8}("__argₛᵧₘ")

function add_arg_to_rewrites!(rewrites::Dict{Any, Any}, lhs::BasicSymbolic{T}) where {T}
    rewrites[lhs] = if get(rewrites, :readable_variables, false)::Bool
        Symbol(string(lhs)::String)
    else
        get_numbered_symbol_with_prefix(ARGNAME_PREFIX, hash(lhs))
    end
end

function (cs::CodegenState{T})(l::Let) where {T}
    if all(x->x isa Assignment && !(x.lhs isa DestructuredArgs), l.pairs)
        dargs = Vector{Assignment}(l.pairs)
    else
        dargs = Assignment[]
        for x in l.pairs
            handle_let_pair!(dargs, x, cs)
        end
    end

    if l.let_block
        scs, bm = enter_scope(cs)
        for asg in dargs
            lhs = asg.lhs
            if lhs isa BasicSymbolic{T} && iscall(lhs)
                add_arg_to_rewrites!(scs.rewrites, lhs)
            end
            scs(asg)
        end
        bindings = exit_scope!(scs, bm)

        scs, bm = enter_scope(cs)
        scs(l.body)
        body = exit_scope!(scs, bm)
        result = Expr(:let, bindings, body)
        return declare!(cs, get_misc_identifier(cs), result)
    end
    isempty(dargs) && return cs(l.body)

    for asg in dargs
        lhs = asg.lhs
        if lhs isa BasicSymbolic{T} && iscall(lhs)
            add_arg_to_rewrites!(cs.rewrites, lhs)
        end
        cs(asg)
    end

    return cs(l.body)
end

function (cs::CodegenState{T})(fn::Func) where {T}
    dargs = Union{Assignment, DestructuredArgs}[]
    for arg in fn.args
        if arg isa BasicSymbolic{T}
            add_arg_to_rewrites!(cs.rewrites, arg)
        elseif arg isa Assignment
            lhs = arg.lhs
            if lhs isa BasicSymbolic{T}
                add_arg_to_rewrites!(cs.rewrites, lhs)
            end
        elseif arg isa DestructuredArgs
            push!(dargs, arg)
        end
    end
    for kw in fn.kwargs
        lhs = kw.lhs
        if lhs isa BasicSymbolic{T}
            # Does not use `add_arg_to_rewrites!` because kwarg names matter
            cs.rewrites[lhs] = Symbol(string(lhs)::String)
        end
    end

    scs, bm = enter_scope(cs)
    scs(Let(dargs, fn.body, false))
    body = exit_scope!(scs, bm)

    fn_args = Expr(:tuple)
    if !isempty(fn.kwargs)
        fn_kws = Expr(:parameters)
        for asg in fn.kwargs
            lhs = manual_dispatch_toexpr(asg.lhs, NameState(cs.rewrites))
            asg_cs, bm = enter_scope(cs)
            asg_cs(asg.rhs)
            rhs = exit_scope!(asg_cs, bm)
            push!(fn_kws.args, Expr(:kw, lhs, rhs))
        end
        push!(fn_args.args, fn_kws)
    end

    for arg in fn.args
        if arg isa BasicSymbolic{T}
            push!(fn_args.args, cs.rewrites[arg]::Symbol)
        elseif arg isa Assignment
            asg_cs, bm = enter_scope(cs)
            asg_cs(arg.rhs)
            rhs = exit_scope!(asg_cs, bm)
            lhs = manual_dispatch_toexpr(arg.lhs, NameState(cs.rewrites))
            push!(fn_args.args, Expr(:(=), lhs, rhs))
        elseif arg isa DestructuredArgs
            name = arg.name
            if name isa BasicSymbolic{T}
                push!(fn_args.args, manual_dispatch_toexpr(name, NameState(cs.rewrites)))
            else
                push!(fn_args.args, name::Symbol)
            end
        else
            push!(fn_args.args, manual_dispatch_toexpr(arg, NameState(cs.rewrites)))
        end
    end

    fn_body = Expr(:block)
    append!(fn_body.args, fn.pre)
    push!(fn_body.args, body)
    return declare!(cs, get_misc_identifier(cs), Expr(:function, fn_args, fn_body))
end

function (cs::CodegenState)(ex::SetArray)
    arr = cs(ex.arr)::Symbol
    lhss = []
    rhss = Symbol[]
    for (i, elem) in enumerate(ex.elems)
        if elem isa AtIndex
            push!(lhss, cs(elem.i))
            push!(rhss, cs(elem.elem))
        else
            push!(lhss, i)
            push!(rhss, cs(elem))
        end
    end

    scs, bm = enter_scope(cs)
    for (lhs, rhs) in zip(lhss, rhss)
        declare!(scs, Expr(:ref, arr, lhs), rhs)
    end
    declare!(scs, get_misc_identifier(cs), ex.return_arr ? arr : nothing)
    body = exit_scope!(scs, bm)
    if ex.inbounds
        body = Expr(:macrocall, nameof(var"@inbounds"), LineNumberNode(0), body)
    end
    return declare!(cs, get_misc_identifier(cs), body)
end

function (cs::CodegenState)(ex::MakeArray)
    _similarto = ex.similarto
    if _similarto isa Type
        similarto = cs(_similarto)
    else
        similarto = declare!(cs, get_misc_identifier(cs), Expr(:call, typeof, cs(_similarto)))
    end
    ndim = ndims(ex.elems)::Int
    elT = ex.output_eltype
    result = Expr(
        :call, create_array, similarto, elT, Expr(:call, Expr(:curly, Val, ndim)),
        Expr(:call, Expr(:curly, Val, size(ex.elems)))
    )
    for x in ex.elems
        push!(result.args, cs(x))
    end
    return declare!(cs, get_misc_identifier(cs), result)
end

function (cs::CodegenState{T})(a::MakeSparseArray{<:SparseMatrixCSC}) where {T}
    I, J, V = findnz(a.array)
    m, n = size(a.array)
    result = Expr(:call, sparse, I, J, cs(Const{T}(V)), m, n)
    return declare!(cs, get_misc_identifier(cs), result)
end

function (cs::CodegenState{T})(a::MakeSparseArray{<:SparseVector}) where {T}
    I, V = findnz(a.array)
    n = length(a.array)
    result = Expr(:call, SparseVector, n, I, cs(Const{T}(V)))
    return declare!(cs, get_misc_identifier(cs), result)
end

function (cs::CodegenState)(a::MakeTuple)
    result = Expr(:tuple)
    for ex in a.elems
        push!(result.args, cs(ex))
    end
    return declare!(cs, get_misc_identifier(cs), result)
end

function (cs::CodegenState)(ex::LiteralExpr)
    declare!(cs, get_misc_identifier(cs), recurse_expr(ex.ex, NameState(cs.rewrites)))
end

function (cs::CodegenState)(f::ForLoop)
    st = NameState(cs.rewrites)
    result = Expr(
        :for, Expr(
            :(=), manual_dispatch_toexpr(f.itervar, st), manual_dispatch_toexpr(f.range, st)
        )
    )
    scs, bm = enter_scope(cs)
    scs(f.body)
    body = exit_scope!(scs, bm)
    push!(result.args, body)
    return declare!(cs, get_misc_identifier(cs), result)
end

function codegen_function!(::Type{Let}, cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    nbors = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    n = length(nbors)
    body_idx = nbors[n]

    # Codegen each symAssignment as a side effect in order
    for i in 1:n-1
        cs(cs.ir[nbors[i]])
    end

    return codegen!(cs, expr_idx, cs(cs.ir[body_idx]))
end

function codegen_function!(::Type{Assignment}, cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    nbors = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    lhs_expr = cs.ir[nbors[1]]
    rhs = cs.ir[nbors[2]]

    # Generate the RHS before the LHS in its own scope. If the RHS involves the LHS
    # (e.g. `x ← x + y`) The `x` on the RHS will not generate correctly without this.
    # The scope allows `x + y` to refer to the symbol `x`, and then for the assignment
    # to redefine `x` as the rewrite created below.
    scoped_cs, bm = enter_scope(cs)
    scoped_cs(rhs)
    rhs_result = exit_scope!(scoped_cs, bm)
    # This is intentionally different from normal assignment codegen. That has to be
    # backward compatible, but this comes from `symAssignment` and doesn't. The old
    # version turns `Assignment(x(t), rhs)` into `x(t) = rhs` which is an inline
    # function. `symAssignment` turns it into `var"x(t)" = rhs`, and to declare
    # a function the user can use `symFunc`.
    lhs = add_arg_to_rewrites!(cs.rewrites, lhs_expr)
    return codegen!(cs, expr_idx, Expr(:(=), lhs, rhs_result))
end

function codegen_function!(::Type{Func}, cs::CodegenState{T}, expr::BasicSymbolic{T}, expr_idx::Integer) where {T}
    nbors = Graphs.outneighbors(cs.ir.dependency_graph, expr_idx)
    args_term_idx = nbors[1]
    body_idx = nbors[2]

    args_nbors = Graphs.outneighbors(cs.ir.dependency_graph, args_term_idx)
    n_args = length(args_nbors)
    body = cs.ir[body_idx]

    # We actually want to use a whole now `CodegenState` here. The arguments to this function
    # can be variables already present in the outer code. An expression already generated in `cs`
    # involving these arguments would then point to the outer code instead of something
    # dependent on the argument of this inner closure.
    #
    # Since we're copying the current `rewrites`, arguments passed to the outer function
    # will be available to the closure.
    new_rewrites = copy(cs.rewrites)

    # (name, elems): array param + the element variables to bind inside the body
    dargs_info = Tuple{BasicSymbolic{T}, Vector{BasicSymbolic{T}}}[]

    for i in 1:n_args
        arg = cs.ir[args_nbors[i]]
        @match arg begin
            BSImpl.Sym(;) => begin
                add_arg_to_rewrites!(new_rewrites, arg)
            end
            BSImpl.Term(; f) && if f === DestructuredArgs end => begin
                da_nbors = Graphs.outneighbors(cs.ir.dependency_graph, args_nbors[i])
                name = cs.ir[da_nbors[1]]
                elems = BasicSymbolic{T}[cs.ir[da_nbors[j]] for j in 2:length(da_nbors)]
                add_arg_to_rewrites!(new_rewrites, name)
                push!(dargs_info, (name, elems))
            end
            _ => nothing
        end
    end

    fn_body_block = let block = Expr(:block), expr = block, inner_cs = CodegenState(
            expr, block, cs.ir, new_rewrites
        )
        # Emit destructuring bindings — elem_i = array_param[i] — before the body
        for (name, elems) in dargs_info
            name_sym = inner_cs.rewrites[name]::Symbol
            for (i, elem) in enumerate(elems)
                elem_name = manual_dispatch_toexpr(elem, NameState(inner_cs.rewrites))
                declare!(inner_cs, elem_name, Expr(:ref, name_sym, i))
            end
        end

        inner_cs(body)
        inner_cs.expr
    end

    fn_args = Expr(:tuple)
    for i in 1:n_args
        arg = cs.ir[args_nbors[i]]
        @match arg begin
            BSImpl.Sym(;) => push!(fn_args.args, new_rewrites[arg]::Symbol)
            BSImpl.Term(; f) && if f === DestructuredArgs end => begin
                da_nbors = Graphs.outneighbors(cs.ir.dependency_graph, args_nbors[i])
                name = cs.ir[da_nbors[1]]
                push!(fn_args.args, new_rewrites[name]::Symbol)
            end
            _ => push!(fn_args.args, manual_dispatch_toexpr(arg, NameState(cs.rewrites)))
        end
    end

    result = codegen!(cs, expr_idx, Expr(:function, fn_args, Expr(:block, fn_body_block)))

    return result
end
