// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license that can be found in the LICENSE file.
module markused

import v.ast
import v.util
import v.pref

// mark_used walks the AST, starting at main() and marks all used fns transitively.
pub fn mark_used(mut table ast.Table, mut pref_ pref.Preferences, ast_files []&ast.File) {
    mut all_fns, all_consts, all_globals, all_decltypes, all_structs := all_global_decl(ast_files)
    mut generic_fns := []&ast.FnDecl{}
    for file in ast_files {
        generic_fns << file.generic_fns
    }
    util.timing_start('MARKUSED')
    defer {
        util.timing_measure('MARKUSED')
    }

    trace_skip_unused := pref_.compile_values['trace_skip_unused'] == 'true'
    trace_skip_unused_all_fns := pref_.compile_values['trace_skip_unused_all_fns'] == 'true'
    trace_skip_unused_fn_names := pref_.compile_values['trace_skip_unused_fn_names'] == 'true'
    trace_skip_unused_just_unused_fns := pref_.compile_values['trace_skip_unused_just_unused_fns'] == 'true'
    used_fns := pref_.compile_values['used_fns']

    string_idx_str := ast.string_type_idx.str()
    array_idx_str := ast.array_type_idx.str()
    map_idx_str := ast.map_type_idx.str()
    ref_map_idx_str := int(ast.map_type.ref()).str()
    ref_densearray_idx_str := int(table.find_type('DenseArray').ref()).str()
    ref_array_idx_str := int(ast.array_type.ref()).str()

    // Functions that must be generated and can't be skipped
    mut all_fn_root_names := []string{}
    mut include_panic_deps := false
    if used_fns != '' {
        aused_fns := used_fns.split(',')
        all_fns_keys := all_fns.keys()
        mut matching := []string{}
        for ufn in aused_fns {
            if ufn.contains('*') {
                matching_fns := all_fns_keys.filter(it.match_glob(ufn))
                if matching_fns.len > 0 {
                    matching << matching_fns
                }
            } else {
                matching << ufn
            }
        }
        all_fn_root_names << matching
        for m in matching {
            println('> used_fn, found matching symbol: ${m}')
        }
    }
    {
        mut core_fns := [
            'main.main',
        ]
        if pref_.is_bare {
            core_fns << 'init_global_allocator' // needed for linux_bare and wasm_bare
        }
        if 'use_libbacktrace' in pref_.compile_defines {
            core_fns << 'print_libbacktrace'
        }
        if 'callstack' in pref_.compile_defines {
            core_fns << ref_array_idx_str + '.push'
            core_fns << ref_array_idx_str + '.pop'
        }
        if pref_.autofree {
            core_fns << string_idx_str + '.clone_static'
            core_fns << string_idx_str + '.option_clone_static'
        }
        if table.used_features.auto_str || pref_.is_shared {
            include_panic_deps = true
            core_fns << 'isnil'
            core_fns << '__new_array'
            core_fns << '__new_array_noscan'
            core_fns << '__new_array_with_multi_default'
            core_fns << '__new_array_with_multi_default_noscan'
            core_fns << '__new_array_with_array_default'
            core_fns << '__new_array_with_array_default_noscan'
            core_fns << 'new_array_from_c_array'
        }
        if table.used_features.arr_prepend {
            core_fns << ref_array_idx_str + '.prepend_many'
            core_fns << ref_array_idx_str + '.prepend_noscan'
        }
        if table.used_features.arr_reverse {
            core_fns << array_idx_str + '.reverse'
        }
        if table.used_features.arr_pop_left {
            core_fns << ref_array_idx_str + '.pop_left'
            core_fns << ref_array_idx_str + '.pop_left_noscan'
        }
        if table.used_features.arr_pop {
            core_fns << ref_array_idx_str + '.pop'
            core_fns << ref_array_idx_str + '.pop_noscan'
        }
        if table.used_features.arr_first {
            core_fns << array_idx_str + '.first'
        }
        if table.used_features.arr_last {
            core_fns << array_idx_str + '.last'
        }
        if table.used_features.arr_insert {
            core_fns << ref_array_idx_str + '.insert_many'
            core_fns << ref_array_idx_str + '.insert_noscan'
        }
        if table.used_features.print_options {
            include_panic_deps = true
            core_fns << '_option_ok'
            core_fns << '_result_ok'
        }
        if table.used_features.anon_fn {
            core_fns << 'memdup_uncollectable'
            core_fns << 'builtin.closure.closure_alloc'
            core_fns << 'builtin.closure.closure_init'
            core_fns << 'builtin.closure.closure_create'
            core_fns << 'builtin.closure.closure_data'
            core_fns << 'builtin.closure.closure_try_destroy'
        }
        if table.used_features.arr_map {
            include_panic_deps = true
            core_fns << '__new_array_with_map_default'
            core_fns << 'new_map_noscan_key'
            core_fns << ref_map_idx_str + '.clone'
            core_fns << ref_densearray_idx_str + '.clone'
            core_fns << map_idx_str + '.clone'
        }
        if pref_.trace_calls || pref_.trace_fns.len > 0 {
            include_panic_deps = true
            core_fns << 'C.gettid'
            core_fns << 'v.trace_calls.on_c_main'
            core_fns << 'v.trace_calls.current_time'
            core_fns << 'v.trace_calls.on_call'
        }
        if 'C.cJSON_Parse' in all_fns {
            core_fns << '_result_ok'
            core_fns << 'tos5'
            core_fns << 'time.unix' // used by json
            core_fns << 'error'
            include_panic_deps = true
        }
        if pref_.should_use_segfault_handler() {
            core_fns << 'v_segmentation_fault_handler'
        }
        if pref_.is_check_overflow {
            // add all fns in `builtin/overflow/overflow.v`
            for op in ['add', 'sub', 'mul'] {
                for typ in ['i8', 'u8', 'i16', 'u16', 'i32', 'u32', 'i64', 'u64'] {
                    core_fns << 'builtin.overflow.${op}_${typ}'
                }
            }
        }
        all_fn_root_names << core_fns
    }
    if pref_.is_bare {
        all_fn_root_names << [
            'strlen',
            'memcmp',
            'memcpy',
            'realloc',
            'vsnprintf',
            'vsprintf',
        ]
    }

    is_noscan_whitelisted := pref_.gc_mode in [.boehm_full_opt, .boehm_incr_opt, .vgc]

    has_noscan := all_fn_root_names.any(it.contains('noscan')
        && it !in ['vcalloc_noscan', 'malloc_noscan'])
    for k, mut mfn in all_fns {
        if trace_skip_unused_all_fns {
            println('k: ${k} | mfn: ${mfn.name}')
        }
        // public/exported functions can not be skipped,
        // especially when producing a shared library:
        if mfn.is_pub && pref_.is_shared {
            all_fn_root_names << k
            continue
        }
        if pref_.translated && mfn.attrs.any(it.name == 'c') {
            all_fn_root_names << k
            continue
        }
        // _noscan functions/methods are selected when the active GC can honor noscan allocations.
        if has_noscan && is_noscan_whitelisted && mfn.name.ends_with('_noscan') {
            all_fn_root_names << k
            continue
        }
        if mfn.is_method {
            method_receiver_typename := table.type_to_str(mfn.receiver.typ)
            if method_receiver_typename == '&sync.Channel' {
                all_fn_root_names << k
                continue
            }
        }
        has_dot := k.contains('.')
        if has_dot {
            if k.ends_with('.init') || k.ends_with('.cleanup') {
                all_fn_root_names << k
                continue
            }
            if pref_.is_prof && (k.starts_with('time.vpc_now') || k.starts_with('v.profile.')) {
                // needed for -profile
                all_fn_root_names << k
                continue
            }
            if k.ends_with('.lock') || k.ends_with('.unlock') || k.ends_with('.rlock')
                || k.ends_with('.runlock') {
                all_fn_root_names << k
                continue
            }
        }

        if k.ends_with('before_request') {
            // TODO: add a more specific check for the .before_request() method in veb apps
            all_fn_root_names << k
            continue
        }
        // testing framework:
        if pref_.is_test {
            if k.starts_with('test_')
                || (has_dot && (k.contains('.test_') || k.contains('.vtest_'))) {
                all_fn_root_names << k
                continue
            }
            if k.starts_with('testsuite_') || (has_dot && k.contains('.testsuite_')) {
                all_fn_root_names << k
                continue
            }
        }
        if pref_.prealloc && k.starts_with('prealloc_') {
            all_fn_root_names << k
            continue
        }
    }

    // handle assertions and testing framework callbacks:
    if pref_.is_debug {
        all_fn_root_names << 'panic_debug'
    }
    // tos3 is used by cgen for typeof() on sum_types at runtime:
    all_fn_root_names << 'tos3'
    if pref_.is_test {
        all_fn_root_names << 'main.cb_assertion_ok'
        all_fn_root_names << 'main.cb_assertion_failed'
        if benched_tests_sym := table.find_sym('main.BenchedTests') {
            bts_type := benched_tests_sym.methods[0].params[0].typ.str()
            all_fn_root_names << bts_type + '.testing_step_start'
            all_fn_root_names << bts_type + '.testing_step_end'
            all_fn_root_names << bts_type + '.end_testing'
            all_fn_root_names << 'main.start_testing'
        }
    }

    handle_veb(mut table, mut all_fn_root_names, 'veb.Result', 'veb.filter', 'veb.Context')

    if 'debug_used_features' in pref_.compile_defines {
        eprintln('> debug_used_features: ${table.used_features}')
    }

    mut walker := Walker.new(
        table:         table
        all_fns:       all_fns
        generic_fns:   generic_fns
        all_consts:    all_consts
        all_globals:   all_globals
        all_decltypes: all_decltypes
        all_structs:   all_structs
        pref:          pref_
        trace_enabled: 'trace_skip_unused_walker' in pref_.compile_defines
    )
    walker.mark_markused_consts() // tagged with `@[markused]`
    walker.mark_markused_globals() // tagged with `@[markused]`
    walker.mark_markused_syms() // tagged with `@[markused]`
    walker.mark_markused_fns() // tagged with `@[markused]`, `@[export]` and veb actions
    walker.mark_markused_decltypes() // tagged with `@[markused]`
    walker.mark_generic_types()

    if pref_.use_cache {
        walker.mark_by_sym_name('IError')
    }

    walker.mark_root_fns(all_fn_root_names)
    walker.mark_generic_fn_instances()

    // Mark all concrete generic type instances as used. These are created by
    // generic_insts_to_concrete() and unwrap_generic_type_ex() for specific
    // type instantiations. The walker may not mark them because it visits
    // generic function bodies with unresolved (generic) AST types.
    for sym in table.type_symbols {
        if sym.info is ast.Struct && sym.info.concrete_types.len > 0 && !sym.info.is_generic {
            walker.mark_by_sym(sym)
        } else if sym.info is ast.SumType && sym.info.concrete_types.len > 0 && !sym.info.is_generic {
            walker.mark_by_sym(sym)
        } else if sym.info is ast.Interface && sym.info.concrete_types.len > 0
            && !sym.info.is_generic {
            walker.mark_by_sym(sym)
        } else if sym.info is ast.GenericInst && sym.info.parent_idx > 0
            && !sym.info.concrete_types.any(it.has_flag(.generic)) {
            walker.mark_by_sym(sym)
        } else if sym.info is ast.Thread && sym.info.return_type != ast.void_type {
            walker.mark_by_sym(sym)
        } else if sym.info is ast.Array {
            elem_sym := table.sym(sym.info.elem_type)
            if elem_sym.info is ast.Thread && elem_sym.info.return_type != ast.void_type {
                walker.mark_by_sym(sym)
            }
        }
    }

    for kcon, con in all_consts {
        if pref_.is_shared && con.is_pub {
            walker.mark_const_as_used(kcon)
            continue
        }
        if pref_.translated && con.attrs.any(it.name == 'export') {
            walker.mark_const_as_used(kcon)
            continue
        }
    }

    if trace_skip_unused_fn_names {
        for key, _ in walker.used_fns {
            println('> used fn key: ${key}')
        }
    }

    walker.finalize(include_panic_deps)

    table.used_features.used_none = walker.used_none
    if walker.used_none == 0 {
        walker.used_fns.delete('${int(ast.none_type)}.str')
    }

    table.used_features.used_fns = walker.used_fns.move()
    table.used_features.used_consts = walker.used_consts.move()
    table.used_features.used_globals = walker.used_globals.move()
    table.used_features.used_syms = walker.used_syms.move()
    table.used_features.used_closures = walker.used_closures

    if trace_skip_unused {
        eprintln('>> t.used_fns: ${table.used_features.used_fns.keys()}')
        eprintln('>> t.used_consts: ${table.used_features.used_consts.keys()}')
        eprintln('>> t.used_globals: ${table.used_features.used_globals.keys()}')
        eprintln('>> t.used_syms: ${table.used_features.used_syms.keys()}')
        eprintln('>> t.used_maps: ${table.used_features.used_maps}')
        eprintln('>> t.used_closures: ${table.used_features.used_closures}')
    }
    if trace_skip_unused_just_unused_fns {
        all_fns_keys := all_fns.keys()
        used_fns_keys := table.used_features.used_fns.keys()
        for k in all_fns_keys {
            if k in used_fns_keys {
                continue
            }
            println('> k: ${k}')
        }
    }
}

fn all_global_decl_in_stmts(stmts []ast.Stmt, mut all_fns map[string]ast.FnDecl, mut all_consts map[string]ast.ConstField, mut all_globals map[string]ast.GlobalField, mut all_decltypes map[string]ast.TypeDecl, mut all_structs map[string]ast.StructDecl) {
    for node in stmts {
        match node {
            ast.FnDecl {
                fkey := node.fkey()
                if fkey !in all_fns || !node.no_body {
                    all_fns[fkey] = node
                }
            }
            ast.ConstDecl {
                for cfield in node.fields {
                    ckey := cfield.name
                    all_consts[ckey] = cfield
                }
            }
            ast.GlobalDecl {
                for gfield in node.fields {
                    gkey := gfield.name
                    all_globals[gkey] = gfield
                }
            }
            ast.StructDecl {
                all_structs[node.name] = node
            }
            ast.TypeDecl {
                if node.is_markused {
                    all_decltypes[node.name] = node
                }
            }
            ast.ExprStmt {
                match node.expr {
                    ast.IfExpr {
                        if node.expr.is_comptime {
                            // top level comptime $if
                            for branch in node.expr.branches {
                                all_global_decl_in_stmts(branch.stmts, mut all_fns, mut all_consts, mut
                                    all_globals, mut all_decltypes, mut all_structs)
                            }
                        }
                    }
                    ast.MatchExpr {
                        if node.expr.is_comptime {
                            // top level comptime $match
                            for branch in node.expr.branches {
                                all_global_decl_in_stmts(branch.stmts, mut all_fns, mut all_consts, mut
                                    all_globals, mut all_decltypes, mut all_structs)
                            }
                        }
                    }
                    else {}
                }
            }
            else {}
        }
    }
}

fn all_global_decl(ast_files []&ast.File) (map[string]ast.FnDecl, map[string]ast.ConstField, map[string]ast.GlobalField, map[string]ast.TypeDecl, map[string]ast.StructDecl) {
    util.timing_start(@METHOD)
    defer {
        util.timing_measure(@METHOD)
    }
    mut all_fns := map[string]ast.FnDecl{}
    mut all_consts := map[string]ast.ConstField{}
    mut all_globals := map[string]ast.GlobalField{}
    mut all_decltypes := map[string]ast.TypeDecl{}
    mut all_structs := map[string]ast.StructDecl{}
    for i in 0 .. ast_files.len {
        all_global_decl_in_stmts(ast_files[i].stmts, mut all_fns, mut all_consts, mut all_globals, mut
            all_decltypes, mut all_structs)
    }
    return all_fns, all_consts, all_globals, all_decltypes, all_structs
}

fn mark_all_methods_used(mut table ast.Table, mut all_fn_root_names []string, typ ast.Type) {
    sym := table.sym(typ)
    styp := int(typ).str()
    for method in sym.methods {
        all_fn_root_names << styp + '.' + method.name
    }
}

fn handle_veb(mut table ast.Table, mut all_fn_root_names []string, result_name string, filter_name string,
    context_name string) {
    // handle veb magic router methods:
    result_type_idx := table.find_type(result_name)
    if result_type_idx == 0 {
        return
    }
    all_fn_root_names << filter_name
    typ_veb_context := table.find_type(context_name).set_nr_muls(1)
    mark_all_methods_used(mut table, mut all_fn_root_names, typ_veb_context)
    for vgt in table.used_features.used_veb_types {
        sym_app := table.sym(vgt)
        pvgt := int(vgt.set_nr_muls(1)).str()
        for m in sym_app.methods {
            mut skip := true
            if m.name == 'before_request' {
                // TODO: handle expansion of method calls in generic functions in a more universal way
                skip = false
            }
            if m.return_type == result_type_idx {
                skip = false
            }
            if skip {
                continue
            }
            // eprintln('vgt: ${vgt} | pvgt: ${pvgt} | sym_app.name: ${sym_app.name} | m.name: ${m.name}')
            all_fn_root_names << pvgt + '.' + m.name
        }
    }
}