/*
 * MRustC - Rust Compiler
 * - By John Hodge (Mutabah/thePowersGang)
 *
 * hir_expand/annotate_value_usage.cpp
 * - Marks _Variable, _Index, _Deref, and _Field nodes with how the result is used
 */
#include <hir/visitor.hpp>
#include <hir/expr.hpp>
#include <hir_typeck/static.hpp>
#include <algorithm>
#include "main_bindings.hpp"
#include <hir/expr_state.hpp>

namespace {

    class ExprVisitor_Mark:
        public ::HIR::ExprVisitor//Def
    {
        struct ClosureScope
        {
            ::HIR::ExprNode_Closure&    node;
            ::std::vector<unsigned int> local_vars;
            // - Lists captured variables to be stored in autogenerated struct (and how they're used, influencing the borrow type)
            ::std::vector<HIR::ExprNode_Closure::AvuCache::Capture> captured_vars;

            ClosureScope(::HIR::ExprNode_Closure& node):
                node(node)
            {
            }
        };

        /// Scope state for generators
        struct CoroutineScope
        {
            static const unsigned STACK_MARKER_LOOP;
            TAGGED_UNION(CRNode, Generator,
                (Generator, ::HIR::ExprNode_Generator*),
                (Async, ::HIR::ExprNode_AsyncBlock*)
            );
            CRNode node;

            // Note: Counts the total number of `yield`s encountered
            // - Loops are pre-counted to either be `-2` (no inner yields) or `-1` (has an inner yield)
            std::vector<unsigned>    yield_stack;
            struct Var {
                // If
                // > `defined_stack` is empty
                // > OR `defined_stack` isn't a prefix of `last_used_stack`
                // > OR `last_used_stack`'s tail contains a yielding loop
                // - Then the variable needs to be saved
                std::vector<unsigned>   defined_stack;
                std::vector<unsigned>   last_used_stack;
                ::HIR::ValueUsage   usage;
            };
            std::map<unsigned, Var> used_variables;

            CoroutineScope(::HIR::ExprNode_Generator& node):
                node(&node)
            {
                // This should always be non-empty
                yield_stack.push_back(0);
            }
            CoroutineScope(::HIR::ExprNode_AsyncBlock& node):
                node(&node)
            {
                // This should always be non-empty
                yield_stack.push_back(0);
            }
        };
        TAGGED_UNION(Scope, None,
            (None, struct{}),
            (Closure, ClosureScope),
            (Coroutine, CoroutineScope)
        );

        const StaticTraitResolve&    m_resolve;
        const ::std::vector< ::HIR::TypeRef>& m_variable_types;

        ::std::vector< ::HIR::ValueUsage>   m_usage;
        ::std::vector< Scope> m_closure_stack;
        /// Disable the capture logic in `ExprNode_Variable` (as the variable was captured using a field access)
        bool    m_ignore_variable_capture;


        struct UsageGuard
        {
            ExprVisitor_Mark& m_parent;
            bool    m_pop;
            UsageGuard(ExprVisitor_Mark& parent, bool pop):
                m_parent(parent),
                m_pop(pop)
            {
            }
            ~UsageGuard()
            {
                if(m_pop) {
                    m_parent.m_usage.pop_back();
                }
            }
        };

        ::HIR::ValueUsage get_usage() const {
            return (m_usage.empty() ? ::HIR::ValueUsage::Move : m_usage.back());
        }
        UsageGuard push_usage(::HIR::ValueUsage u) {
            if( get_usage() == u ) {
                return UsageGuard(*this, false);
            }
            else {
                m_usage.push_back( u );
                return UsageGuard(*this, true);
            }
        }

    public:
        ExprVisitor_Mark(const StaticTraitResolve& resolve, const ::std::vector< ::HIR::TypeRef>& variable_types):
            m_resolve(resolve)
            , m_variable_types(variable_types)
            , m_ignore_variable_capture(false)
        {}

        void visit_root(::HIR::ExprPtr& root_ptr)
        {
            // HACK: Pre-visit all nodes to find closures, and mark those as !Copy
            // - Closures are marked as Copy until this pass, where they need to default to not-Copy so usage defaults to the most restrictive
            {
                struct IV: public ::HIR::ExprVisitorDef {
                    void visit(::HIR::ExprNode_Closure& node) override {
                        node.m_is_copy = false;
                        ::HIR::ExprVisitorDef::visit(node);
                    }
                }   iv;
                root_ptr->visit(iv);
            }

            assert(root_ptr);
            root_ptr->m_usage = this->get_usage();
            auto expected_size = m_usage.size();
            root_ptr->visit( *this );
            assert( m_usage.size() == expected_size );
        }
        void visit_node_ptr(::HIR::ExprNodeP& node_ptr) override
        {
            assert(node_ptr);

            const auto& node_ref = *node_ptr;
            const char* node_tyname = typeid(node_ref).name();
            TRACE_FUNCTION_FR(&*node_ptr << " " << node_tyname << ": " << this->get_usage(), node_ptr->m_usage);

            node_ptr->m_usage = this->get_usage();

            auto expected_size = m_usage.size();
            node_ptr->visit( *this );
            assert( m_usage.size() == expected_size );
        }

        void visit(::HIR::ExprNode_Block& node) override
        {
            CoroutineScope* scope = m_closure_stack.size() > 0 ? m_closure_stack.back().opt_Coroutine() : nullptr;
            if(scope) {
                scope->yield_stack.push_back(0);
            }

            auto _ = this->push_usage( ::HIR::ValueUsage::Move );

            for( auto& subnode : node.m_nodes ) {
                this->visit_node_ptr(subnode);
            }
            if( node.m_value_node )
                this->visit_node_ptr(node.m_value_node);

            if(scope) {
                // NOTE: Have to get the pointer again (it might have changed due to inner push/pop)
                scope = &m_closure_stack.back().as_Coroutine();
                scope->yield_stack.pop_back();
            }
        }
        void visit(::HIR::ExprNode_ConstBlock& node) override
        {
            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr(node.m_inner);
        }

        void visit(::HIR::ExprNode_Asm& node) override
        {
            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            for(auto& v : node.m_outputs)
            {
                this->visit_node_ptr(v.value);
            }
            for(auto& v : node.m_inputs)
            {
                this->visit_node_ptr(v.value);
            }
        }
        void visit(::HIR::ExprNode_Asm2& node) override
        {
            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            for(auto& v : node.m_params)
            {
                TU_MATCH_HDRA( (v), { )
                TU_ARMA(Const, e) {
                    visit_node_ptr(e);
                    }
                TU_ARMA(Sym, e) {
                    }
                TU_ARMA(RegSingle, e) {
                    visit_node_ptr(e.val);
                    }
                TU_ARMA(Reg, e) {
                    if(e.val_in)    visit_node_ptr(e.val_in);
                    if(e.val_out)   visit_node_ptr(e.val_out);
                    }
                }
            }
        }
        void visit(::HIR::ExprNode_Return& node) override
        {
            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr( node.m_value );
        }
        void visit(::HIR::ExprNode_Yield& node) override
        {
            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr( node.m_value );

            // `yield`: Increase all stack entries that don't correspond to a loop
            CoroutineScope* scope = m_closure_stack.size() > 0 ? m_closure_stack.back().opt_Coroutine() : nullptr;
            if(scope && scope->node.is_Generator())
            {
                for(size_t i = 0; i < scope->yield_stack.size(); i ++)
                {
                    if(scope->yield_stack[i] != CoroutineScope::STACK_MARKER_LOOP) {
                        // Assert that adding 1 won't overflow (... that'd be >65k yield statements)
                        assert(scope->yield_stack[i] < CoroutineScope::STACK_MARKER_LOOP-1);
                        scope->yield_stack[i] += 1;
                    }
                }
                DEBUG("yield stack=[" << scope->yield_stack << "]");
            }
        }
        void visit(::HIR::ExprNode_AWait& node) override
        {
            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr( node.m_value );

            // `yield`: Increase all stack entries that don't correspond to a loop
            CoroutineScope* scope = m_closure_stack.size() > 0 ? m_closure_stack.back().opt_Coroutine() : nullptr;
            if(scope && scope->node.is_Async())
            {
                for(size_t i = 0; i < scope->yield_stack.size(); i ++)
                {
                    if(scope->yield_stack[i] != CoroutineScope::STACK_MARKER_LOOP) {
                        // Assert that adding 1 won't overflow (... that'd be >65k yield statements)
                        assert(scope->yield_stack[i] < CoroutineScope::STACK_MARKER_LOOP-1);
                        scope->yield_stack[i] += 1;
                    }
                }
                DEBUG("await stack=[" << scope->yield_stack << "]");
            }
        }
        void visit(::HIR::ExprNode_Let& node) override
        {
            add_defs_from_pattern(node.span(), node.m_pattern);
            if( node.m_value )
            {
                auto _ = this->push_usage( this->get_usage_for_pattern(node.span(), node.m_pattern, node.m_type) );
                this->visit_node_ptr( node.m_value );
            }
        }
        void visit(::HIR::ExprNode_Loop& node) override
        {
            CoroutineScope* scope = m_closure_stack.size() > 0 ? m_closure_stack.back().opt_Coroutine() : nullptr;
            if(scope) {
                // 1. Determine if there's a yield within this
                struct InnerVisitor: public ::HIR::ExprVisitorDef
                {
                    bool m_has_yield = false;

                    void visit(::HIR::ExprNode_Closure& ) override
                    {
                        // Ignore inner closures (yields can't pass them)
                    }
                    void visit(::HIR::ExprNode_Generator& ) override
                    {
                        // Ignore inner generators (yields can't pass them)
                    }
                    void visit(::HIR::ExprNode_AsyncBlock& ) override
                    {
                        // Ignore inner asyncs (yields can't pass them)
                    }
                    void visit(::HIR::ExprNode_Yield& node) override
                    {
                        m_has_yield = true;
                        // Don't bother recursing, we've found a yield
                    }
                } v;
                v.visit_node_ptr(node.m_code);

                // 2. If so, push `STACK_MARKER_LOOP`
                if( v.m_has_yield ) {
                    DEBUG("Loop with inner yield");
                    scope->yield_stack.push_back(CoroutineScope::STACK_MARKER_LOOP);
                }
                else {
                    // Clear `scope` to prevent the pop
                    DEBUG("Loop without inner yield");
                    scope = nullptr;
                }
            }

            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr( node.m_code );

            if(scope) {
                // NOTE: Have to get the pointer again (it might have changed due to inner push/pop)
                scope = &m_closure_stack.back().as_Coroutine();
                scope->yield_stack.pop_back();
            }
        }
        void visit(::HIR::ExprNode_LoopControl& node) override
        {
            // NOTE: Leaf
            if( node.m_value )
            {
                this->visit_node_ptr(node.m_value);
            }
        }
        void visit(::HIR::ExprNode_Match& node) override
        {
            {
                const auto& val_ty = node.m_value->m_res_type;
                ::HIR::ValueUsage   vu = ::HIR::ValueUsage::Unknown;
                for( const auto& arm : node.m_arms )
                {
                    for( const auto& pat : arm.m_patterns ) {
                        DEBUG("_Match: " << pat);
                        vu = ::std::max( vu, this->get_usage_for_pattern(node.span(), pat, val_ty) );
                    }
                }
                if( vu == ::HIR::ValueUsage::Unknown ) {
                    DEBUG("No value usage for pattern arms (no arms?), default to borrow");
                    vu = ::HIR::ValueUsage::Borrow;
                }
                auto _ = this->push_usage( vu );
                this->visit_node_ptr( node.m_value );
            }

            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            for(auto& arm : node.m_arms)
            {
                for( const auto& pat : arm.m_patterns )
                    add_defs_from_pattern(node.span(), pat);
                for( auto& c : arm.m_guards) {
                    auto _ = this->push_usage( this->get_usage_for_pattern(c.val->span(), c.pat, c.val->m_res_type) );
                    this->visit_node_ptr( c.val );
                    add_defs_from_pattern(node.span(), c.pat);
                }
                this->visit_node_ptr( arm.m_code );
            }
        }
        void visit(::HIR::ExprNode_If& node) override
        {
            auto _ = this->push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr( node.m_cond );
            this->visit_node_ptr( node.m_true );
            if( node.m_false ) {
                this->visit_node_ptr( node.m_false );
            }
        }

        void visit(::HIR::ExprNode_Assign& node) override
        {
            {
                auto _ = this->push_usage( ::HIR::ValueUsage::Mutate );
                this->visit_node_ptr(node.m_slot);
            }
            {
                auto _ = this->push_usage( ::HIR::ValueUsage::Move );
                this->visit_node_ptr(node.m_value);
            }
        }
        void visit(::HIR::ExprNode_UniOp& node) override
        {
            m_usage.push_back( ::HIR::ValueUsage::Move );

            this->visit_node_ptr(node.m_value);

            m_usage.pop_back();
        }
        void visit(::HIR::ExprNode_Borrow& node) override
        {
            switch(node.m_type)
            {
            case ::HIR::BorrowType::Shared:
                m_usage.push_back( ::HIR::ValueUsage::Borrow );
                break;
            case ::HIR::BorrowType::Unique:
                m_usage.push_back( ::HIR::ValueUsage::Mutate );
                break;
            case ::HIR::BorrowType::Owned:
                m_usage.push_back( ::HIR::ValueUsage::Move );
                break;
            }

            this->visit_node_ptr(node.m_value);

            m_usage.pop_back();
        }
        void visit(::HIR::ExprNode_RawBorrow& node) override
        {
            switch(node.m_type)
            {
            case ::HIR::BorrowType::Shared:
                m_usage.push_back( ::HIR::ValueUsage::Borrow );
                break;
            case ::HIR::BorrowType::Unique:
                m_usage.push_back( ::HIR::ValueUsage::Mutate );
                break;
            case ::HIR::BorrowType::Owned:
                m_usage.push_back( ::HIR::ValueUsage::Move );
                break;
            }

            this->visit_node_ptr(node.m_value);

            m_usage.pop_back();
        }

        void visit(::HIR::ExprNode_BinOp& node) override
        {
            switch(node.m_op)
            {
            case ::HIR::ExprNode_BinOp::Op::CmpEqu:
            case ::HIR::ExprNode_BinOp::Op::CmpNEqu:
            case ::HIR::ExprNode_BinOp::Op::CmpLt:
            case ::HIR::ExprNode_BinOp::Op::CmpLtE:
            case ::HIR::ExprNode_BinOp::Op::CmpGt:
            case ::HIR::ExprNode_BinOp::Op::CmpGtE:
                m_usage.push_back( ::HIR::ValueUsage::Borrow );
                break;
            default:
                m_usage.push_back( ::HIR::ValueUsage::Move );
                break;
            }

            this->visit_node_ptr(node.m_left);
            this->visit_node_ptr(node.m_right);

            m_usage.pop_back();
        }
        void visit(::HIR::ExprNode_Cast& node) override
        {
            auto _ = push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr(node.m_value);
        }
        void visit(::HIR::ExprNode_Unsize& node) override
        {
            // TODO: Why does Unsize have a usage of Move?
            //auto _ = push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr(node.m_value);
        }
        void visit(::HIR::ExprNode_Index& node) override
        {
            // TODO: Override to ::Borrow if Res: Copy and moving
            if( this->get_usage() == ::HIR::ValueUsage::Move && m_resolve.type_is_copy(node.span(), node.m_res_type) ) {
                auto _ = push_usage( ::HIR::ValueUsage::Borrow );
                this->visit_node_ptr(node.m_value);
            }
            else {
                this->visit_node_ptr(node.m_value);
            }

            auto _ = push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr(node.m_index);
        }
        void visit(::HIR::ExprNode_Deref& node) override
        {
            if( this->get_usage() == ::HIR::ValueUsage::Move && m_resolve.type_is_copy(node.span(), node.m_res_type) ) {
                auto _ = push_usage( ::HIR::ValueUsage::Borrow );
                this->visit_node_ptr(node.m_value);
            }
            // Pointers only need a borrow to be derefernced.
            else if( node.m_value->m_res_type.data().is_Pointer() ) {
                auto _ = push_usage( ::HIR::ValueUsage::Borrow );
                this->visit_node_ptr(node.m_value);
            }
            else {
                this->visit_node_ptr(node.m_value);
            }
        }

        void visit(::HIR::ExprNode_Emplace& node) override
        {
            if( node.m_type == ::HIR::ExprNode_Emplace::Type::Noop ) {
                if( node.m_place )
                    this->visit_node_ptr(node.m_place);
                this->visit_node_ptr(node.m_value);
            }
            else {
                auto _ = push_usage( ::HIR::ValueUsage::Move );
                if( node.m_place )
                    this->visit_node_ptr(node.m_place);
                this->visit_node_ptr(node.m_value);
            }
        }

        void visit(::HIR::ExprNode_Field& node) override
        {
            bool is_copy = m_resolve.type_is_copy(node.span(), node.m_res_type);
            DEBUG("ty = " << node.m_res_type << ", is_copy=" << is_copy);

            bool saved_ignore_variable_capure = m_ignore_variable_capture;
            // Rust 2021 - Closures (and I assume generators) capture fields separately
            if(m_resolve.m_crate.m_edition >= AST::Edition::Rust2021) {
                // Is there an active closure (or generator) AND this node isn't in a chain that has already been considered.
                if( !m_closure_stack.empty() && !m_ignore_variable_capture ) {
                    ::std::vector<RcString> fields;
                    fields.push_back(node.m_field);

                    // Extract a chain of field names
                    auto* inner = node.m_value.get();
                    while( auto* inner_field = dynamic_cast<::HIR::ExprNode_Field*>(inner) ) {
                        fields.push_back(inner_field->m_field);
                        inner = inner_field->m_value.get();
                    }
                    if( auto* inner_deref = dynamic_cast<::HIR::ExprNode_Deref*>(inner) ) {
                        if( inner_deref->m_value->m_res_type.data().is_Borrow() ) {
                            fields.push_back(RcString());
                            inner = inner_deref->m_value.get();
                        }
                    }
                    // and if the final value is a variable, then insert into the captures
                    if( auto* inner_var = dynamic_cast<::HIR::ExprNode_Variable*>(inner) ) {
                        std::reverse(fields.begin(), fields.end());

                        mark_used_variable(node.span(), inner_var->m_slot, fields, this->get_usage());

                        m_ignore_variable_capture = true;
                    }
                }
            }

            // If taking this field by value, but the type is Copy - pretend it's a borrow.
            if( this->get_usage() == ::HIR::ValueUsage::Move && is_copy ) {
                auto _ = push_usage( ::HIR::ValueUsage::Borrow );
                this->visit_node_ptr(node.m_value);
            }
            else {
                this->visit_node_ptr(node.m_value);
            }

            m_ignore_variable_capture = saved_ignore_variable_capure;
        }

        void visit(::HIR::ExprNode_TupleVariant& node) override
        {
            auto _ = push_usage( ::HIR::ValueUsage::Move );

            for( auto& val : node.m_args )
                this->visit_node_ptr(val);
        }
        void visit(::HIR::ExprNode_CallPath& node) override
        {
            auto _ = push_usage( ::HIR::ValueUsage::Move );

            for( auto& val : node.m_args )
                this->visit_node_ptr(val);
        }
        void visit(::HIR::ExprNode_CallValue& node) override
        {
            // TODO: Different usage based on trait.
            ::HIR::ValueUsage   vu = ::HIR::ValueUsage::Borrow;

            if( const auto* node_pp = TU_OPT1(node.m_value->m_res_type.data(), NodeType, .opt_Closure()) ) {
                assert(*node_pp);
                if( (*node_pp)->m_class == ::HIR::ExprNode_Closure::Class::Unknown ) {
                    auto _ = push_usage( ::HIR::ValueUsage::Move );
                    this->visit_node_ptr(node.m_value);
                }
                switch((*node_pp)->m_class)
                {
                case ::HIR::ExprNode_Closure::Class::Unknown:
                    BUG(node.span(), "CallValue with unknown closure class - " << node.m_value->m_res_type);
                    break;
                case ::HIR::ExprNode_Closure::Class::NoCapture:
                case ::HIR::ExprNode_Closure::Class::Shared:
                    node.m_trait_used = ::HIR::ExprNode_CallValue::TraitUsed::Fn;
                    break;
                case ::HIR::ExprNode_Closure::Class::Mut:
                    node.m_trait_used = ::HIR::ExprNode_CallValue::TraitUsed::FnMut;
                    break;
                case ::HIR::ExprNode_Closure::Class::Once:
                    node.m_trait_used = ::HIR::ExprNode_CallValue::TraitUsed::FnOnce;
                    break;
                }
            }

            switch( node.m_trait_used )
            {
            case ::HIR::ExprNode_CallValue::TraitUsed::Unknown:
                BUG(node.span(), "Annotate usage when CallValue trait is unknown");
                break;
            case ::HIR::ExprNode_CallValue::TraitUsed::Fn:
                vu = ::HIR::ValueUsage::Borrow;
                break;
            case ::HIR::ExprNode_CallValue::TraitUsed::FnMut:
                vu = ::HIR::ValueUsage::Mutate;
                break;
            case ::HIR::ExprNode_CallValue::TraitUsed::FnOnce:
                vu = ::HIR::ValueUsage::Move;
                break;
            }
            {
                auto _ = push_usage( vu );
                this->visit_node_ptr(node.m_value);
            }

            auto _ = push_usage( ::HIR::ValueUsage::Move );
            for( auto& val : node.m_args )
                this->visit_node_ptr(val);
        }
        void visit(::HIR::ExprNode_CallMethod& node) override
        {
            {
                assert(node.m_cache.m_fcn);
                ::HIR::ValueUsage   vu = ::HIR::ValueUsage::Borrow;
                switch(node.m_cache.m_fcn->m_receiver)
                {
                case ::HIR::Function::Receiver::Free:
                    BUG(node.span(), "_CallMethod resolved to free function");
                case ::HIR::Function::Receiver::Value:
                case ::HIR::Function::Receiver::Box:
                case ::HIR::Function::Receiver::Custom:
                case ::HIR::Function::Receiver::BorrowOwned:
                    vu = ::HIR::ValueUsage::Move;
                    break;
                case ::HIR::Function::Receiver::BorrowUnique:
                    vu = ::HIR::ValueUsage::Mutate;
                    break;
                case ::HIR::Function::Receiver::BorrowShared:
                    vu = ::HIR::ValueUsage::Borrow;
                    break;
                //case ::HIR::Function::Receiver::PointerMut:
                //case ::HIR::Function::Receiver::PointerConst:
                }
                auto _ = push_usage( vu );
                this->visit_node_ptr(node.m_value);
            }
            auto _ = push_usage( ::HIR::ValueUsage::Move );
            for( auto& val : node.m_args )
                this->visit_node_ptr(val);
        }

        void visit(::HIR::ExprNode_Literal& node) override
        {
        }
        void visit(::HIR::ExprNode_UnitVariant& node) override
        {
        }
        void visit(::HIR::ExprNode_PathValue& node) override
        {
        }
        void visit(::HIR::ExprNode_Variable& node) override
        {
            DEBUG("_Variable: #" << node.m_slot << " '" << node.m_name << "' " << node.m_usage);
            if( !m_closure_stack.empty() && !m_ignore_variable_capture )
            {
                mark_used_variable(node.span(), node.m_slot, {}, node.m_usage);
            }
        }
        void visit(::HIR::ExprNode_ConstParam& node) override
        {
        }

        void visit(::HIR::ExprNode_StructLiteral& node) override
        {
            const auto& sp = node.span();
            const auto& ty_path = node.m_real_path;
            if( node.m_base_value ) {
                bool is_moved = false;
                const auto& tpb = node.m_base_value->m_res_type.data().as_Path().binding;
                const ::HIR::t_struct_fields* fields_ptr;
                if( tpb.is_Enum() ) {
                    const auto& enm = *tpb.as_Enum();
                    auto idx = enm.find_variant(ty_path.m_path.components().back());
                    ASSERT_BUG(sp, idx != SIZE_MAX, "");
                    const auto& var_ty = enm.m_data.as_Data()[idx].type;
                    const auto& str = *var_ty.data().as_Path().binding.as_Struct();
                    ASSERT_BUG(sp, str.m_data.is_Named(), "");
                    fields_ptr = &str.m_data.as_Named();
                }
                else if( tpb.is_Union() ) {
                    fields_ptr = &tpb.as_Union()->m_variants;
                }
                else {
                    const auto& str = *tpb.as_Struct();
                    ASSERT_BUG(sp, str.m_data.is_Named(), "");
                    fields_ptr = &str.m_data.as_Named();
                }
                const auto& fields = *fields_ptr;

                ::std::vector<bool> provided_mask( fields.size() );
                for( const auto& fld : node.m_values ) {
                    unsigned idx = ::std::find_if( fields.begin(), fields.end(), [&](const HIR::StructField& x){ return x.name == fld.first; }) - fields.begin();
                    provided_mask[idx] = true;
                }

                const auto monomorph_cb = MonomorphStatePtr(nullptr, &ty_path.m_params, nullptr);
                for( unsigned int i = 0; i < fields.size(); i ++ ) {
                    if( ! provided_mask[i] ) {
                        const auto& ty_o = fields[i].ty;
                        ::HIR::TypeRef  tmp;
                        const auto& ty_m = monomorphise_type_with_opt(node.span(), tmp, ty_o, monomorph_cb);
                        bool is_copy = m_resolve.type_is_copy(node.span(), ty_m);
                        if( !is_copy ) {
                            DEBUG("- Field " << i << " " << fields[i].name << ": " << ty_m << " moved");
                            is_moved = true;
                        }
                    }
                }

                // If only Copy fields will be used, set usage to Borrow
                auto _ = push_usage( is_moved ? ::HIR::ValueUsage::Move : ::HIR::ValueUsage::Borrow );
                this->visit_node_ptr(node.m_base_value);
            }

            auto _ = push_usage( ::HIR::ValueUsage::Move );
            for( auto& fld_val : node.m_values ) {
                this->visit_node_ptr(fld_val.second);
            }
        }
        void visit(::HIR::ExprNode_Tuple& node) override
        {
            auto _ = push_usage( ::HIR::ValueUsage::Move );
            for( auto& val : node.m_vals ) {
                this->visit_node_ptr(val);
            }
        }
        void visit(::HIR::ExprNode_ArrayList& node) override
        {
            auto _ = push_usage( ::HIR::ValueUsage::Move );
            for( auto& val : node.m_vals ) {
                this->visit_node_ptr(val);
            }
        }
        void visit(::HIR::ExprNode_ArraySized& node) override
        {
            auto _ = push_usage( ::HIR::ValueUsage::Move );
            this->visit_node_ptr(node.m_val);
        }

        void visit(::HIR::ExprNode_Closure& node) override
        {
            if(!node.m_code)
            {
                DEBUG("Already expanded (via consteval?)");
                return ;
            }

            m_closure_stack.push_back(ClosureScope(node));

            for(const auto& arg : node.m_args)
                add_defs_from_pattern(node.span(), arg.first);

            if(node.m_code)
            {
                auto _ = push_usage( ::HIR::ValueUsage::Move );
                this->visit_node_ptr(node.m_code);
            }

            auto scope = std::move(m_closure_stack.back().as_Closure());
            m_closure_stack.pop_back();

            if(node.m_class == ::HIR::ExprNode_Closure::Class::Unknown) {
                DEBUG("> Class is still `Unknown`, set to `NoCapture`");
                node.m_class = ::HIR::ExprNode_Closure::Class::NoCapture;
            }
            else {
                DEBUG("= Class is " << static_cast<int>(node.m_class));
            }

            if( node.m_is_move )
            {
                DEBUG("> Tagged with `move` - upgrading all usage to `Move`");
                for(auto& cap : scope.captured_vars)
                {
                    if( cap.fields.size() > 0 && cap.fields[0] == "" ) {
                        // Derefs stay as-is?
                        // TODO: Only clear if the value isn't `Copy`
                        cap.fields.clear();
                    }
                    cap.usage = ::HIR::ValueUsage::Move;
                }
            }
            else
            {
                for(auto& cap : scope.captured_vars)
                {
                    if( cap.usage == ::HIR::ValueUsage::Borrow ) {
                        ::HIR::TypeRef  tmp_ty;
                        const auto* cap_ty_p = &m_variable_types.at(cap.root_slot);
                        for(const auto& n : cap.fields) {
                            tmp_ty = m_resolve.get_field_type(node.span(), *cap_ty_p, n);
                            m_resolve.expand_associated_types(node.span(), tmp_ty);
                            cap_ty_p = &tmp_ty;
                        }
                        if( cap_ty_p->data().is_Borrow() && cap_ty_p->data().as_Borrow().type == ::HIR::BorrowType::Shared ) {
                            DEBUG("> Upgrade capture " << cap.root_slot << cap.fields << " to Move, as it's a shared borrow");
                            cap.usage = ::HIR::ValueUsage::Move;
                        }
                    }
                }
            }

            // Determine if the closure is `Copy`
            {
                node.m_is_copy = true;
                for(auto& cap : scope.captured_vars)
                {
                    switch(cap.usage)
                    {
                    case HIR::ValueUsage::Unknown:
                        BUG(node.span(), "Usage of capture #" << cap.root_slot << cap.fields << " is unknown");
                    case HIR::ValueUsage::Borrow:
                        break;
                    case HIR::ValueUsage::Mutate:
                        node.m_is_copy = false;
                        break;
                    case HIR::ValueUsage::Move: {
                        const auto* ty = &m_variable_types.at(cap.root_slot);
                        HIR::TypeRef    tmp_ty;
                        for(const auto& fld : cap.fields) {
                            tmp_ty = m_resolve.get_field_type(node.span(), *ty, fld);
                            ty = &tmp_ty;
                        }
                        if( !m_resolve.type_is_copy(node.span(), *ty) )
                        {
                            node.m_is_copy = false;
                        }
                        break; }
                    }
                }
                if( node.m_is_copy )
                {
                    DEBUG("> Copy closure");
                }
            }

            if(!m_closure_stack.empty())
            {
                DEBUG("> Apply to parent");
                for(const auto& v : scope.captured_vars)
                {
                    mark_used_variable(node.span(), v.root_slot, v.fields, v.usage);
                }
            }

            node.m_avu_cache.captured_vars = std::move(scope.captured_vars);
            node.m_avu_cache.local_vars = std::move(scope.local_vars);
        }
        void visit(::HIR::ExprNode_Generator& node) override
        {
            m_closure_stack.push_back(CoroutineScope(node));

            {
                auto _ = push_usage( ::HIR::ValueUsage::Move );
                this->visit_node_ptr(node.m_code);
            }

            auto ent = std::move(m_closure_stack.back().as_Coroutine());
            m_closure_stack.pop_back();

            apply_coroutine(node.span(), node.m_is_move, node.m_avu_cache, ent.used_variables);
        }
        void visit(::HIR::ExprNode_GeneratorWrapper& node) override
        {
            BUG(node.span(), "");
        }

        void visit(::HIR::ExprNode_AsyncBlock& node) override
        {
            // NOTE: This is a near-clone of `ExprNode_Generator` (they're very similar, and at one point
            // futures/async were directly implemented as syntax sugar around generators)
            m_closure_stack.push_back(CoroutineScope(node));

            {
                auto _ = push_usage( ::HIR::ValueUsage::Move );
                this->visit_node_ptr(node.m_code);
            }

            auto scope = std::move(m_closure_stack.back().as_Coroutine());
            m_closure_stack.pop_back();

            apply_coroutine(node.span(), node.m_is_move, node.m_avu_cache, scope.used_variables);
        }

    private:
        void add_var_def_closure(const Span& sp, ClosureScope& e, unsigned int slot)
        {
            auto it = ::std::lower_bound(e.local_vars.begin(), e.local_vars.end(), slot);
            if( it == e.local_vars.end() || *it != slot ) {
                e.local_vars.insert(it, slot);
            }
        }
        void add_var_def_generator(const Span& sp, CoroutineScope& scope, unsigned int slot)
        {
            auto& e = scope.used_variables[slot];
            DEBUG("_Variable: #" << slot << " '?' stack=[" << scope.yield_stack << "]");
            e.defined_stack = scope.yield_stack;
        }

        void add_var_def(const Span& sp, unsigned int slot)
        {
            assert(m_closure_stack.size() > 0);
            auto& ent = m_closure_stack.back();
            TU_MATCH_HDRA( (ent), {)
            TU_ARMA(None, e) throw "";
            TU_ARMA(Closure, e) {
                add_var_def_closure(sp, e, slot);
                }
            TU_ARMA(Coroutine, e) {
                add_var_def_generator(sp, e, slot);
                }
            }
        }
        void add_closure_def_from_pattern(const Span& sp, const ::HIR::Pattern& pat)
        {
            // Add binding indexes to m_closure_defs
            for(const auto& pb : pat.m_bindings ) {
                add_var_def(sp, pb.m_slot);
            }

            // Recurse
            TU_MATCH_HDRA((pat.m_data), {)
            TU_ARMA(Any, e) {
                }
            TU_ARMA(Value, e) {
                }
            TU_ARMA(Range, e) {
                }
            TU_ARMA(Box, e) {
                add_closure_def_from_pattern(sp, *e.sub);
                }
            TU_ARMA(Ref, e) {
                add_closure_def_from_pattern(sp, *e.sub);
                }
            TU_ARMA(Tuple, e) {
                for( const auto& subpat : e.sub_patterns )
                    add_closure_def_from_pattern(sp, subpat);
                }
            TU_ARMA(SplitTuple, e) {
                for( const auto& subpat : e.leading )
                    add_closure_def_from_pattern(sp, subpat);
                for( const auto& subpat : e.trailing )
                    add_closure_def_from_pattern(sp, subpat);
                }
            TU_ARMA(Slice, e) {
                for(const auto& sub : e.sub_patterns)
                    add_closure_def_from_pattern(sp, sub);
                }
            TU_ARMA(SplitSlice, e) {
                for(const auto& sub : e.leading)
                    add_closure_def_from_pattern( sp, sub );
                for(const auto& sub : e.trailing)
                    add_closure_def_from_pattern( sp, sub );
                if( e.extra_bind.is_valid() ) {
                    add_var_def(sp, e.extra_bind.m_slot);
                }
                }

            // - Enums/Structs
            TU_ARMA(PathValue, e) {
                }
            TU_ARMA(PathTuple, e) {
                for(const auto& field : e.leading) {
                    add_closure_def_from_pattern(sp, field);
                }
                for(const auto& field : e.trailing) {
                    add_closure_def_from_pattern(sp, field);
                }
                }
            TU_ARMA(PathNamed, e) {
                for( auto& field_pat : e.sub_patterns ) {
                    add_closure_def_from_pattern(sp, field_pat.second);
                }
                }
            
            TU_ARMA(Or, e) {
                assert(e.size() > 0);
                add_closure_def_from_pattern(sp, e.front());
                }
            }
        }
        void add_defs_from_pattern(const Span& sp, const ::HIR::Pattern& pat)
        {
            if(!m_closure_stack.empty()) {
                add_closure_def_from_pattern(sp, pat);
            }
        }

        ::HIR::ValueUsage get_usage_for_pattern_binding(const Span& sp, const ::HIR::PatternBinding& pb, const ::HIR::TypeRef& ty) const
        {
            switch( pb.m_type )
            {
            case ::HIR::PatternBinding::Type::Move:
                if( m_resolve.type_is_copy(sp, ty) )
                    return ::HIR::ValueUsage::Borrow;
                else
                    return ::HIR::ValueUsage::Move;
            case ::HIR::PatternBinding::Type::MutRef:
                return ::HIR::ValueUsage::Mutate;
            case ::HIR::PatternBinding::Type::Ref:
                return ::HIR::ValueUsage::Borrow;
            }
            throw "";
        }

        ::HIR::ValueUsage get_usage_for_pattern(const Span& sp, const ::HIR::Pattern& pat, const ::HIR::TypeRef& outer_ty) const
        {
            if( pat.m_bindings.size() > 0 )
            {
                auto vu = ::HIR::ValueUsage::Borrow;
                for(const auto& pb : pat.m_bindings ) {
                    vu = std::max(vu, get_usage_for_pattern_binding(sp, pb, outer_ty));
                }
                return vu;
            }

            // Implicit derefs
            const ::HIR::TypeRef* typ = &outer_ty;
            for(size_t i = 0; i < pat.m_implicit_deref_count; i ++)
            {
                typ = &typ->data().as_Borrow().inner;
            }
            const ::HIR::TypeRef& ty = *typ;

            TU_MATCH_HDRA( (pat.m_data), {)
            TU_ARMA(Any, pe) {
                return ::HIR::ValueUsage::Borrow;
                }
            TU_ARMA(Box, pe) {
                // NOTE: Specific to `owned_box`
                const auto& sty = ty.data().as_Path().path.m_data.as_Generic().m_params.m_types.at(0);
                return get_usage_for_pattern(sp, *pe.sub, sty);
                }
            TU_ARMA(Ref, pe) {
                return get_usage_for_pattern(sp, *pe.sub, ty.data().as_Borrow().inner);
                }
            TU_ARMA(Tuple, pe) {
                ASSERT_BUG(sp, ty.data().is_Tuple(), "Tuple pattern with non-tuple type - " << ty);
                const auto& subtys = ty.data().as_Tuple();
                assert(pe.sub_patterns.size() == subtys.size());
                auto rv = ::HIR::ValueUsage::Borrow;
                for(unsigned int i = 0; i < subtys.size(); i ++)
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pe.sub_patterns[i], subtys[i]));
                return rv;
                }
            TU_ARMA(SplitTuple, pe) {
                ASSERT_BUG(sp, ty.data().is_Tuple(), "SplitTuple pattern with non-tuple type - " << ty);
                const auto& subtys = ty.data().as_Tuple();
                assert(pe.leading.size() + pe.trailing.size() <= subtys.size());
                auto rv = ::HIR::ValueUsage::Borrow;
                for(unsigned int i = 0; i < pe.leading.size(); i ++)
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pe.leading[i], subtys[i]));
                for(unsigned int i = 0; i < pe.trailing.size(); i ++)
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pe.trailing[i], subtys[subtys.size() - pe.trailing.size() + i]));
                return rv;
                }
            TU_ARMA(PathValue, pe) {
                return ::HIR::ValueUsage::Borrow;
                }
            TU_ARMA(PathTuple, pe) {
                assert(!pe.binding.is_Unbound());

                const auto& flds = ::HIR::pattern_get_tuple(sp, pe.path, pe.binding);
                if(pe.is_split) {
                    assert(pe.leading.size() + pe.trailing.size() <= flds.size());
                }
                else {
                    assert(pe.leading.size() == flds.size());
                    assert(pe.trailing.size() == 0);
                }

                // TODO: Is it possible to avoid monomorphising here?
                assert(pe.path.m_data.is_Generic());
                auto monomorph_state = MonomorphStatePtr(nullptr,  &pe.path.m_data.as_Generic().m_params, nullptr);

                auto rv = ::HIR::ValueUsage::Borrow;
                for(unsigned int i = 0; i < pe.leading.size(); i ++)
                {
                    auto sty = m_resolve.monomorph_expand(sp, flds[i].ent, monomorph_state);
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pe.leading[i], sty));
                }
                for(unsigned int i = 0; i < pe.trailing.size(); i ++)
                {
                    auto sty = m_resolve.monomorph_expand(sp, flds[flds.size() - pe.trailing.size() + i].ent, monomorph_state);
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pe.trailing[i], sty));
                }
                return rv;
                }
            TU_ARMA(PathNamed, pe) {
                assert(!pe.binding.is_Unbound());

                if( pe.is_wildcard() )
                    return ::HIR::ValueUsage::Borrow;
                if( pe.sub_patterns.empty() )
                    return ::HIR::ValueUsage::Borrow;

                const auto& flds = ::HIR::pattern_get_named(sp, pe.path, pe.binding);
                auto monomorph_state = MonomorphStatePtr(nullptr,  &pe.path.m_data.as_Generic().m_params, nullptr);

                auto rv = ::HIR::ValueUsage::Borrow;
                for(const auto& fld_pat : pe.sub_patterns)
                {
                    auto fld_it = ::std::find_if(flds.begin(), flds.end(), [&](const HIR::StructField& x){return x.name == fld_pat.first;});
                    ASSERT_BUG(sp, fld_it != flds.end(), "Unable to find field " << fld_pat.first);

                    auto sty = m_resolve.monomorph_expand(sp, fld_it->ty, monomorph_state);
                    rv = ::std::max(rv, get_usage_for_pattern(sp, fld_pat.second, sty));
                }
                return rv;
                }
            TU_ARMA(Value, pe) {
                return ::HIR::ValueUsage::Borrow;
                }
            TU_ARMA(Range, pe) {
                return ::HIR::ValueUsage::Borrow;
                }
            TU_ARMA(Slice, pe) {
                const auto& inner_ty = (ty.data().is_Array() ? ty.data().as_Array().inner : ty.data().as_Slice().inner);
                auto rv = ::HIR::ValueUsage::Borrow;
                for(const auto& pat : pe.sub_patterns)
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pat, inner_ty));
                return rv;
                }
            TU_ARMA(SplitSlice, pe) {
                const auto& inner_ty = (ty.data().is_Array() ? ty.data().as_Array().inner : ty.data().as_Slice().inner);
                auto rv = ::HIR::ValueUsage::Borrow;
                for(const auto& pat : pe.leading)
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pat, inner_ty));
                for(const auto& pat : pe.trailing)
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pat, inner_ty));
                if( pe.extra_bind.is_valid() )
                    rv = ::std::max(rv, get_usage_for_pattern_binding(sp, pe.extra_bind, inner_ty));
                return rv;
                }
            TU_ARMA(Or, pe) {
                auto rv = ::HIR::ValueUsage::Borrow;
                for(const auto& pat : pe)
                    rv = ::std::max(rv, get_usage_for_pattern(sp, pat, ty));
                return rv;
                }
            }
            throw "";
        }

        // --- Closure/generator usage tracking ---
        /// Update usage depending on the type
        ::HIR::ValueUsage get_real_usage(const Span& sp, unsigned slot, const std::vector<RcString>& fields, ::HIR::ValueUsage usage)
        {
            // If the usage is move, update it depending on the type
            if( usage == ::HIR::ValueUsage::Move )
            {
                const auto* ty = &m_variable_types.at(slot);
                HIR::TypeRef    tmp_ty;
                for(const auto& name : fields) {
                    tmp_ty = m_resolve.get_field_type(sp, *ty, name);
                    ty = &tmp_ty;
                }

                // Copy types just need a borrow
                if( m_resolve.type_is_copy(sp, *ty) ) {
                    usage = ::HIR::ValueUsage::Borrow;
                }
                // `&mut` types get re-borrowed
                // - The reborrow pass is AFTER this pass
                else if( ty->data().is_Borrow() && ty->data().as_Borrow().type == ::HIR::BorrowType::Unique ) {
                    usage = ::HIR::ValueUsage::Mutate;
                }
                else {
                }
            }
            return usage;
        }
        void mark_used_variable_closure(const Span& sp, ClosureScope& closure_rec, unsigned slot, std::vector<RcString> fields, ::HIR::ValueUsage usage)
        {
            DEBUG("(" << slot << ", [" << fields << "], usage=" << usage << ")");
            const auto& closure_defs = closure_rec.local_vars;
            auto& closure = closure_rec.node;

            if( ::std::binary_search(closure_defs.begin(), closure_defs.end(), slot) ) {
                // Ignore, this is local to the current closure
                return ;
            }
            usage = get_real_usage(sp, slot, fields, usage);

            HIR::ExprNode_Closure::AvuCache::Capture    new_ent;
            new_ent.root_slot = slot;
            new_ent.fields = std::move(fields);
            new_ent.usage = usage;

            // `equal_range` considering one being a prefix of the other as equal
            auto its = ::std::equal_range(closure_rec.captured_vars.begin(), closure_rec.captured_vars.end(), new_ent,
                [](const HIR::ExprNode_Closure::AvuCache::Capture& a, const HIR::ExprNode_Closure::AvuCache::Capture& b)->bool {
                    if( a.root_slot < b.root_slot ) {
                        return true;
                    }
                    else if( a.root_slot > b.root_slot ) {
                        return false;
                    }
                    else {  // a.root_slot == b.root_slot

                        if( a.fields == b.fields ) {
                            return false;
                        }
                        auto prefix_len = std::min(a.fields.size(), b.fields.size());
                        for(size_t i = 0; i < prefix_len; i ++) {
                            if( b.fields[i] != a.fields[i] ) {
                                return b.fields[i] < a.fields[i];
                            }
                        }
                        // Equal up to the prefix.
                        return false;
                    }
                });
            if( its.first == its.second ) {
                // Doesn't exist yet
                DEBUG("Insert");
                closure_rec.captured_vars.insert(its.first, new_ent);
            }
            else if( its.first->fields.size() <= new_ent.fields.size() ) {
                DEBUG("new longer");
                assert(its.first->root_slot == new_ent.root_slot);
                for(size_t i = 0; i < its.first->fields.size(); i ++) {
                    assert(its.first->fields[i] == new_ent.fields[i]);
                }
                ASSERT_BUG(sp, its.first + 1 == its.second, "Prefix total match, but multiple matching entries?");
                its.first->usage = std::max(its.first->usage, new_ent.usage);
            }
            else {
                DEBUG("new shorter or equal");
                assert(its.first->root_slot == new_ent.root_slot);
                assert(its.first->fields.size() >= new_ent.fields.size());
                for(size_t i = 0; i < new_ent.fields.size(); i ++) {
                    assert(its.first->fields[i] == new_ent.fields[i]);
                }

                if( its.first + 1 == its.second ) {
                    // Single pre-existing entry
                }
                else {
                    // Multiple pre-existing entries
                    assert(its.first->fields.size() > new_ent.fields.size());
                    for(auto it = its.first + 1; it != its.second; ++it) {
                        its.first->usage = std::max(its.first->usage, it->usage);
                    }
                    closure_rec.captured_vars.erase(its.first + 1, its.second);
                }
                its.first->usage = std::max(its.first->usage, new_ent.usage);
                if( its.first->fields.size() != new_ent.fields.size() ) {
                    its.first->fields = new_ent.fields;
                }
            }

            const char* cap_type_name = "?";
            switch( usage )
            {
            case ::HIR::ValueUsage::Unknown:
                BUG(sp, "Unknown usage of variable " << slot);
            case ::HIR::ValueUsage::Borrow:
                cap_type_name = "Borrow";
                closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Shared);
                break;
            case ::HIR::ValueUsage::Mutate:
                cap_type_name = "Mutate";
                closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Mut);
                break;
            case ::HIR::ValueUsage::Move:
                // TODO: Why is this disabled? Maybe for efficiency? as copies are marked as ValueUsage::Borrow anyway
                if( m_resolve.type_is_copy( sp, m_variable_types.at(slot) ) ) {
                    cap_type_name = "Borrow (copy)";
                    closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Shared);
                }
                else {
                    cap_type_name = "Move";
                    closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Once);
                }
                break;
            }
            DEBUG("Captured " << slot << " - " << m_variable_types.at(slot) << " :: " << cap_type_name);
        }
        void mark_used_variable_generator(const Span& sp, CoroutineScope& scope, unsigned int slot, std::vector<RcString> fields, ::HIR::ValueUsage usage)
        {
            auto& e = scope.used_variables[slot];
            e.last_used_stack = scope.yield_stack;
            e.usage = std::max(e.usage, get_real_usage(sp, slot, fields, usage));
            DEBUG("Used #" << slot << " :: stack=[" << e.last_used_stack << "]");
        }
        void mark_used_variable(const Span& sp, unsigned int slot, std::vector<RcString> fields, ::HIR::ValueUsage usage)
        {
            assert(m_closure_stack.size() > 0);
            auto& ent = m_closure_stack.back();
            TU_MATCH_HDRA( (ent), {)
            TU_ARMA(None, e) throw "";
            TU_ARMA(Closure, e) {
                mark_used_variable_closure(sp, e, slot, fields, usage);
                }
            TU_ARMA(Coroutine, e) {
                mark_used_variable_generator(sp, e, slot, fields, usage);
                }
            }
        }

        void apply_coroutine(const Span& sp, bool is_move, ::HIR::ExprNode_Generator::AvuCache& avu_cache, std::map<unsigned, CoroutineScope::Var>& used_variables)
        {
            // If this closure is a move closure, mutate `captured_vars` such that all captures are tagged with ValueUsage::Move
            if( is_move )
            {
                for(auto& cap : used_variables)
                {
                    if(cap.second.defined_stack.empty())
                    {
                        cap.second.usage = ::HIR::ValueUsage::Move;
                    }
                }
            }
            else
            {
                for(auto& cap : used_variables)
                {
                    if( cap.second.usage == ::HIR::ValueUsage::Borrow ) {
                        ::HIR::TypeRef  tmp_ty;
                        const auto* cap_ty_p = &m_variable_types.at(cap.first);
                        //for(const auto& n : cap.fields) {
                        //    tmp_ty = m_resolve.get_field_type(node.span(), *cap_ty_p, n);
                        //    m_resolve.expand_associated_types(node.span(), tmp_ty);
                        //    cap_ty_p = &tmp_ty;
                        //}
                        if( cap_ty_p->data().is_Borrow() && cap_ty_p->data().as_Borrow().type == ::HIR::BorrowType::Shared ) {
                            DEBUG("> Upgrade capture #" << cap.first<< " to Move, as it's a shared borrow");
                            cap.second.usage = ::HIR::ValueUsage::Move;
                        }
                    }
                }
            }

            // - Apply into the parent stack
            if( m_closure_stack.size() > 0 )
            {
                DEBUG("> Apply to parent");
                for(const auto& cap : used_variables)
                {
                    // Only apply already-defined variables (ones without a set `defined_stack`)
                    if(cap.second.defined_stack.empty())
                    {
                        mark_used_variable(sp, cap.first, {}, cap.second.usage);
                    }
                }
            }

            // Optimisation: Pre-reserve the space required
            {
                size_t n_caps = 0;
                size_t n_locals = 0;
                for(const auto& cap : used_variables)
                {
                    if( cap.second.defined_stack.empty() ) {
                        n_caps += 1;
                    }
                    else {
                        n_locals += 1;
                    }
                }
                avu_cache.captured_vars.reserve(n_caps);
                avu_cache.local_vars.reserve(n_locals);
            }
            for(const auto& cap : used_variables)
            {
                if( cap.second.defined_stack.empty() ) {
                    avu_cache.captured_vars.push_back(std::make_pair(cap.first, cap.second.usage));
                }
                else {
                    avu_cache.local_vars.push_back(cap.first);
                }
            }
        }
    };
    const unsigned ExprVisitor_Mark::CoroutineScope::STACK_MARKER_LOOP = ~0u;


    class OuterVisitor:
        public ::HIR::Visitor
    {
        StaticTraitResolve   m_resolve;
    public:
        OuterVisitor(const ::HIR::Crate& crate):
            m_resolve(crate)
        {}

        void visit_expr(::HIR::ExprPtr& exp) override {
            if( exp )
            {
                ExprVisitor_Mark    ev { m_resolve, exp.m_bindings };
                ev.visit_root( exp );
            }
        }

        // ------
        // Code-containing items
        // ------
        void visit_function(::HIR::ItemPath p, ::HIR::Function& item) override {
            auto _ = this->m_resolve.set_item_generics(item.m_params);
            DEBUG("Function " << p);
            ::HIR::Visitor::visit_function(p, item);
        }
        void visit_static(::HIR::ItemPath p, ::HIR::Static& item) override {
            // NOTE: No generics
            ::HIR::Visitor::visit_static(p, item);
        }
        void visit_constant(::HIR::ItemPath p, ::HIR::Constant& item) override {
            // NOTE: No generics
            ::HIR::Visitor::visit_constant(p, item);
        }
        void visit_enum(::HIR::ItemPath p, ::HIR::Enum& item) override {
            auto _ = this->m_resolve.set_item_generics(item.m_params);
            ::HIR::Visitor::visit_enum(p, item);
        }


        void visit_trait(::HIR::ItemPath p, ::HIR::Trait& item) override {
            auto _ = this->m_resolve.set_impl_generics(MetadataType::TraitObject, item.m_params);
            ::HIR::Visitor::visit_trait(p, item);
        }

        void visit_type_impl(::HIR::TypeImpl& impl) override
        {
            TRACE_FUNCTION_F("impl " << impl.m_type);
            auto _ = this->m_resolve.set_impl_generics(impl.m_type, impl.m_params);

            ::HIR::Visitor::visit_type_impl(impl);
        }
        void visit_trait_impl(const ::HIR::SimplePath& trait_path, ::HIR::TraitImpl& impl) override
        {
            TRACE_FUNCTION_F("impl " << trait_path << " for " << impl.m_type);
            auto _ = this->m_resolve.set_impl_generics(impl.m_type, impl.m_params);

            ::HIR::Visitor::visit_trait_impl(trait_path, impl);
        }
    };
}

void HIR_Expand_AnnotateUsage_Expr(const ::HIR::Crate& crate, const ::HIR::ItemPath& ip, ::HIR::ExprPtr& exp)
{
    TRACE_FUNCTION_F(ip);
    assert(exp);
    StaticTraitResolve   resolve { crate };
    resolve.set_both_generics_raw(exp.m_state->m_impl_generics, exp.m_state->m_item_generics);
    ExprVisitor_Mark    ev { resolve, exp.m_bindings };
    ev.visit_root(exp);
}

void HIR_Expand_AnnotateUsage(::HIR::Crate& crate)
{
    OuterVisitor    ov(crate);
    ov.visit_crate( crate );
}
