// See the file "COPYING" in the main distribution directory for copyright.

// Classes for run-time initialization and management of C++ values used
// by the generated code.

// See InitsInfo.h for a discussion of initialization issues and the
// associated strategies for dealing with them.

#include "zeek/Expr.h"
#include "zeek/module_util.h"
#include "zeek/script_opt/CPP/RuntimeInitSupport.h"

#pragma once

namespace zeek::detail {

using FileValPtr = IntrusivePtr<FileVal>;
using FuncValPtr = IntrusivePtr<FuncVal>;

class InitsManager;

// Helper function that takes a (large) array of int's and from them
// constructs the corresponding vector-of-vector-of-indices.  Each
// vector-of-indices is represented first by an int specifying its
// size, and then that many int's for its values.  We recognize the
// end of the array upon encountering a "size" entry of END_OF_VEC_VEC.
//
// Returns how many elements were processed out of "inits", including its
// terminator.
extern size_t generate_indices_set(int* inits, std::vector<std::vector<int>>& indices_set);

// The same but for one more level of vector construction. The source array
// has sub-arrays terminated with END_OF_VEC_VEC per the above, and the whole
// shebang is terminated with END_OF_VEC_VEC_VEC.
//
// Returns the vector construction.
extern std::vector<std::vector<std::vector<int>>> generate_indices_set(int* inits);

// These need to be distinct from any values that can appear, which means
// they should be negative, and not -1, which is used as a "N/A" value.
#define END_OF_VEC_VEC -100
#define END_OF_VEC_VEC_VEC -200

// An abstract helper class used to access elements of an initialization vector.
// We need the abstraction because InitsManager below needs to be able to refer
// to any of a range of templated classes.
class CPP_AbstractInitAccessor {
public:
    virtual ~CPP_AbstractInitAccessor() {}
    virtual ValPtr Get(int index) const { return nullptr; }
};

// Convenient way to refer to an offset associated with a particular Zeek type.
// A "struct" rather than a std::pair because C++ compilers are terribly slow
// at initializing large numbers of the latter.
struct CPP_ValElem {
    TypeTag tag;
    int offset;
};

// This class groups together all of the vectors needed for run-time
// initialization.  We gather them together into a single object so as
// to avoid wiring in a set of globals that the various initialization
// methods have to know about.
class InitsManager {
public:
    InitsManager(std::vector<CPP_ValElem>& _const_vals,
                 std::map<TypeTag, std::shared_ptr<CPP_AbstractInitAccessor>>& _consts,
                 std::vector<std::vector<int>>& _indices, std::vector<const char*>& _strings,
                 std::vector<p_hash_type>& _hashes, std::vector<TypePtr>& _types,
                 std::vector<AttributesPtr>& _attributes, std::vector<AttrPtr>& _attrs,
                 std::vector<CallExprPtr>& _call_exprs)
        : const_vals(_const_vals),
          consts(_consts),
          indices(_indices),
          strings(_strings),
          hashes(_hashes),
          types(_types),
          attributes(_attributes),
          attrs(_attrs),
          call_exprs(_call_exprs) {}

    // Provides generic access to Zeek constant values based on a single
    // index.
    ValPtr ConstVals(int offset) const {
        auto& cv = const_vals[offset];
        return Consts(cv.tag, cv.offset);
    }

    // Retrieves the Zeek constant value for a particular Zeek type.
    ValPtr Consts(TypeTag tag, int index) const { return consts[tag]->Get(index); }

    // Accessors for the sundry initialization vectors, each retrieving
    // a specific element identified by an index/offset.
    const std::vector<int>& Indices(int offset) const { return indices[offset]; }
    const char* Strings(int offset) const {
        ASSERT(offset >= 0 && offset < static_cast<int>(strings.size()));
        ASSERT(strings[offset]);
        return strings[offset];
    }
    const p_hash_type Hashes(int offset) const {
        ASSERT(offset >= 0 && offset < static_cast<int>(hashes.size()));
        return hashes[offset];
    }
    const TypePtr& Types(int offset) const {
        ASSERT(offset >= 0 && offset < static_cast<int>(types.size()));
        ASSERT(types[offset]);
        return types[offset];
    }
    const AttributesPtr& Attributes(int offset) const {
        ASSERT(offset >= 0 && offset < static_cast<int>(attributes.size()));
        ASSERT(attributes[offset]);
        return attributes[offset];
    }
    const AttrPtr& Attrs(int offset) const {
        ASSERT(offset >= 0 && offset < static_cast<int>(attrs.size()));
        ASSERT(attrs[offset]);
        return attrs[offset];
    }
    const CallExprPtr& CallExprs(int offset) const {
        ASSERT(offset >= 0 && offset < static_cast<int>(call_exprs.size()));
        ASSERT(call_exprs[offset]);
        return call_exprs[offset];
    }

private:
    std::vector<CPP_ValElem>& const_vals;
    std::map<TypeTag, std::shared_ptr<CPP_AbstractInitAccessor>>& consts;
    std::vector<std::vector<int>>& indices;
    std::vector<const char*>& strings;
    std::vector<p_hash_type>& hashes;
    std::vector<TypePtr>& types;
    std::vector<AttributesPtr>& attributes;
    std::vector<AttrPtr>& attrs;
    std::vector<CallExprPtr>& call_exprs;
};

// Manages an initialization vector of the given type.
template<class T>
class CPP_Init {
public:
    virtual ~CPP_Init() {}

    // Pre-initializes the given element of the vector, if necessary.
    virtual void PreInit(InitsManager* im, std::vector<T>& inits_vec, int offset) const {}

    // Initializes the given element of the vector.
    virtual void Generate(InitsManager* im, std::vector<T>& inits_vec, int offset) const {}
};

// Abstract class for creating a collection of initializers.  T1 is
// the type of the generated vector, T2 the type of its initializers.
template<class T1, class T2>
class CPP_AbstractInits {
public:
    CPP_AbstractInits(std::vector<T1>& _inits_vec, int _offsets_set, std::vector<T2> _inits)
        : inits_vec(_inits_vec), offsets_set(_offsets_set), inits(std::move(_inits)) {
        // Compute how big to make the vector.
        int num_inits = 0;

        for ( const auto& cohort : inits )
            num_inits += cohort.size();

        inits_vec.resize(num_inits);
    }

    // Initialize the given cohort of elements.
    void InitializeCohort(InitsManager* im, int cohort) {
        // Get this object's vector-of-vector-of-indices.
        auto& offsets_vec = im->Indices(offsets_set);

        if ( cohort == 0 )
            DoPreInits(im, offsets_vec);

        // Get the vector-of-indices for this cohort.
        auto& cohort_offsets = im->Indices(offsets_vec[cohort]);

        InitializeCohortWithOffsets(im, cohort, cohort_offsets);
    }

protected:
    virtual void InitializeCohortWithOffsets(InitsManager* im, int cohort, const std::vector<int>& cohort_offsets) {}

    // Pre-initialize all elements requiring it.
    virtual void DoPreInits(InitsManager* im, const std::vector<int>& offsets_vec) {}

    // The initialization vector in its entirety.
    std::vector<T1>& inits_vec;

    // A meta-index for retrieving the vector-of-vector-of-indices.
    int offsets_set;

    // Indexed by cohort.
    std::vector<T2> inits;
};

// Manages an initialization vector that uses "custom" initializers
// (tailored ones rather than initializers based on indexing).
template<class T>
using CPP_InitVec = std::vector<std::shared_ptr<CPP_Init<T>>>;
template<class T>
class CPP_CustomInits : public CPP_AbstractInits<T, CPP_InitVec<T>> {
public:
    CPP_CustomInits(std::vector<T>& _inits_vec, int _offsets_set, std::vector<CPP_InitVec<T>> _inits)
        : CPP_AbstractInits<T, CPP_InitVec<T>>(_inits_vec, _offsets_set, std::move(_inits)) {}

private:
    void DoPreInits(InitsManager* im, const std::vector<int>& offsets_vec) override {
        int cohort = 0;
        for ( const auto& co : this->inits ) {
            auto& cohort_offsets = im->Indices(offsets_vec[cohort]);
            for ( auto i = 0U; i < co.size(); ++i )
                co[i]->PreInit(im, this->inits_vec, cohort_offsets[i]);
            ++cohort;
        }
    }

    void InitializeCohortWithOffsets(InitsManager* im, int cohort, const std::vector<int>& cohort_offsets) override {
        // Loop over the cohort's elements to initialize them.
        auto& co = this->inits[cohort];
        for ( auto i = 0U; i < co.size(); ++i )
            co[i]->Generate(im, this->inits_vec, cohort_offsets[i]);
    }
};

// Provides access to elements of an initialization vector of the given type.
template<class T>
class CPP_InitAccessor : public CPP_AbstractInitAccessor {
public:
    CPP_InitAccessor(std::vector<T>& _inits_vec) : inits_vec(_inits_vec) {}

    ValPtr Get(int index) const override { return inits_vec[index]; }

private:
    std::vector<T>& inits_vec;
};

// A type used for initializations that are based on indices into
// initialization vectors.
using ValElemVec = std::vector<int>;
using ValElemVecVec = std::vector<ValElemVec>;

// Manages an initialization vector of the given type whose elements are
// built up from previously constructed values in other initialization vectors.
template<class T>
class CPP_IndexedInits : public CPP_AbstractInits<T, ValElemVecVec> {
public:
    CPP_IndexedInits(std::vector<T>& _inits_vec, int _offsets_set, int* raw_inits)
        : CPP_AbstractInits<T, ValElemVecVec>(_inits_vec, _offsets_set, generate_indices_set(raw_inits)) {}

protected:
    void InitializeCohortWithOffsets(InitsManager* im, int cohort, const std::vector<int>& cohort_offsets) override;

    // Note, in the following we pass in the inits_vec ("ivec"), even though
    // the method will have direct access to it, because we want to use
    // overloading to dispatch to custom generation for different types of
    // values.
    void Generate(InitsManager* im, std::vector<EnumValPtr>& ivec, int offset, ValElemVec& init_vals);
    void Generate(InitsManager* im, std::vector<StringValPtr>& ivec, int offset, ValElemVec& init_vals);
    void Generate(InitsManager* im, std::vector<PatternValPtr>& ivec, int offset, ValElemVec& init_vals);
    void Generate(InitsManager* im, std::vector<ListValPtr>& ivec, int offset, ValElemVec& init_vals) const;
    void Generate(InitsManager* im, std::vector<VectorValPtr>& ivec, int offset, ValElemVec& init_vals) const;
    void Generate(InitsManager* im, std::vector<RecordValPtr>& ivec, int offset, ValElemVec& init_vals) const;
    void Generate(InitsManager* im, std::vector<TableValPtr>& ivec, int offset, ValElemVec& init_vals) const;
    void Generate(InitsManager* im, std::vector<FileValPtr>& ivec, int offset, ValElemVec& init_vals) const;
    void Generate(InitsManager* im, std::vector<FuncValPtr>& ivec, int offset, ValElemVec& init_vals) const;
    void Generate(InitsManager* im, std::vector<AttrPtr>& ivec, int offset, ValElemVec& init_vals) const;
    void Generate(InitsManager* im, std::vector<AttributesPtr>& ivec, int offset, ValElemVec& init_vals) const;

    // The TypePtr initialization vector requires special treatment, since
    // it has to dispatch on subclasses of TypePtr.
    virtual void Generate(InitsManager* im, std::vector<TypePtr>& ivec, int offset, ValElemVec& init_vals) const {
        ASSERT(0);
    }
};

// A specialization of CPP_IndexedInits that supports initializing based
// on subclasses of TypePtr.
class CPP_TypeInits : public CPP_IndexedInits<TypePtr> {
public:
    CPP_TypeInits(std::vector<TypePtr>& _inits_vec, int _offsets_set, int* raw_inits)
        : CPP_IndexedInits<TypePtr>(_inits_vec, _offsets_set, raw_inits) {}

protected:
    void DoPreInits(InitsManager* im, const std::vector<int>& offsets_vec) override;
    void PreInit(InitsManager* im, int offset, ValElemVec& init_vals);

    void Generate(InitsManager* im, std::vector<TypePtr>& ivec, int offset, ValElemVec& init_vals) const override;

    TypePtr BuildEnumType(InitsManager* im, ValElemVec& init_vals) const;
    TypePtr BuildOpaqueType(InitsManager* im, ValElemVec& init_vals) const;
    TypePtr BuildTypeType(InitsManager* im, ValElemVec& init_vals) const;
    TypePtr BuildVectorType(InitsManager* im, ValElemVec& init_vals) const;
    TypePtr BuildTypeList(InitsManager* im, ValElemVec& init_vals, int offset) const;
    TypePtr BuildTableType(InitsManager* im, ValElemVec& init_vals, int offset) const;
    TypePtr BuildFuncType(InitsManager* im, ValElemVec& init_vals) const;
    TypePtr BuildRecordType(InitsManager* im, ValElemVec& init_vals, int offset) const;
};

// Abstract class for initializing basic (non-compound) constants.  T1 is
// the Zeek type for the constructed constant, T2 is the C++ type of its
// initializer.
//
// In principle we could derive this from CPP_AbstractInits, though to do so
// we'd need to convert the initializers to a vector-of-vector-of-T2, which
// would trade complexity here for complexity in InitsInfo.  So we instead
// keep this class distinct, since at heart it's a simpler set of methods
// and that way we can keep them as such here.
template<class T1, typename T2>
class CPP_AbstractBasicConsts {
public:
    CPP_AbstractBasicConsts(std::vector<T1>& _inits_vec, int _offsets_set, std::vector<T2> _inits)
        : inits_vec(_inits_vec), offsets_set(_offsets_set), inits(std::move(_inits)) {
        inits_vec.resize(inits.size());
    }

    void InitializeCohort(InitsManager* im, int cohort) {
        ASSERT(cohort == 0);
        auto& offsets_vec = im->Indices(offsets_set);
        auto& cohort_offsets = im->Indices(offsets_vec[cohort]);
        for ( auto i = 0U; i < inits.size(); ++i )
            InitElem(im, cohort_offsets[i], i);
    }

protected:
    virtual void InitElem(InitsManager* im, int offset, int index) { ASSERT(0); }

protected:
    // See CPP_AbstractInits for the nature of these.
    std::vector<T1>& inits_vec;
    int offsets_set;
    std::vector<T2> inits;
};

// Class for initializing a basic constant of Zeek type T1, using initializers
// of C++ type T2.  T1 is an intrusive pointer to a T3 type; for example, if
// T1 is a BoolValPtr then T3 will be BoolVal.
template<class T1, typename T2, class T3>
class CPP_BasicConsts : public CPP_AbstractBasicConsts<T1, T2> {
public:
    CPP_BasicConsts(std::vector<T1>& _inits_vec, int _offsets_set, std::vector<T2> _inits)
        : CPP_AbstractBasicConsts<T1, T2>(_inits_vec, _offsets_set, std::move(_inits)) {}

    void InitElem(InitsManager* /* im */, int offset, int index) override {
        this->inits_vec[offset] = make_intrusive<T3>(this->inits[index]);
    }
};

// Specific classes for basic constants that use string-based constructors.
class CPP_AddrConsts : public CPP_AbstractBasicConsts<AddrValPtr, int> {
public:
    CPP_AddrConsts(std::vector<AddrValPtr>& _inits_vec, int _offsets_set, std::vector<int> _inits)
        : CPP_AbstractBasicConsts<AddrValPtr, int>(_inits_vec, _offsets_set, std::move(_inits)) {}

    void InitElem(InitsManager* im, int offset, int index) override {
        auto s = im->Strings(this->inits[index]);
        this->inits_vec[offset] = make_intrusive<AddrVal>(s);
    }
};

class CPP_SubNetConsts : public CPP_AbstractBasicConsts<SubNetValPtr, int> {
public:
    CPP_SubNetConsts(std::vector<SubNetValPtr>& _inits_vec, int _offsets_set, std::vector<int> _inits)
        : CPP_AbstractBasicConsts<SubNetValPtr, int>(_inits_vec, _offsets_set, std::move(_inits)) {}

    void InitElem(InitsManager* im, int offset, int index) override {
        auto s = im->Strings(this->inits[index]);
        this->inits_vec[offset] = make_intrusive<SubNetVal>(s);
    }
};

// Class for initializing a Zeek global.  These don't go into an initialization
// vector, so we use void* as the underlying type.
class CPP_GlobalInit : public CPP_Init<void*> {
public:
    CPP_GlobalInit(IDPtr& _global, const char* _name, int _type, int _attrs, int _val, bool _exported,
                   bool _func_with_no_val)
        : CPP_Init<void*>(),
          global(_global),
          name(_name),
          type(_type),
          attrs(_attrs),
          val(_val),
          exported(_exported),
          func_with_no_val(_func_with_no_val) {}

    void Generate(InitsManager* im, std::vector<void*>& /* inits_vec */, int /* offset */) const override;

protected:
    IDPtr& global;
    const char* name;
    int type;
    int attrs;
    int val;
    bool exported;
    bool func_with_no_val;
};

// Abstract class for constructing a CallExpr to evaluate a Zeek expression.
class CPP_AbstractCallExprInit : public CPP_Init<CallExprPtr> {
public:
    CPP_AbstractCallExprInit() : CPP_Init<CallExprPtr>() {}
};

// Constructs a CallExpr that calls a given CPPFunc subclass.
template<class T>
class CPP_CallExprInit : public CPP_AbstractCallExprInit {
public:
    CPP_CallExprInit(CallExprPtr& _e_var) : CPP_AbstractCallExprInit(), e_var(_e_var) {}

    void Generate(InitsManager* /* im */, std::vector<CallExprPtr>& inits_vec, int offset) const override {
        auto wrapper_class = make_intrusive<T>();
        auto func_val = make_intrusive<FuncVal>(wrapper_class);
        auto func_expr = make_intrusive<ConstExpr>(func_val);
        auto empty_args = make_intrusive<ListExpr>();

        e_var = make_intrusive<CallExpr>(func_expr, empty_args);
        inits_vec[offset] = e_var;
    }

private:
    // Where to store the expression once we've built it.
    CallExprPtr& e_var;
};

// Abstract class for registering a lambda defined in terms of a CPPStmt.
class CPP_AbstractLambdaRegistration : public CPP_Init<void*> {
public:
    CPP_AbstractLambdaRegistration() : CPP_Init<void*>() {}
};

// Registers a lambda defined in terms of a given CPPStmt subclass.
template<class T>
class CPP_LambdaRegistration : public CPP_AbstractLambdaRegistration {
public:
    CPP_LambdaRegistration(const char* _name, int _func_type, p_hash_type _h, bool _has_captures)
        : CPP_AbstractLambdaRegistration(), name(_name), func_type(_func_type), h(_h), has_captures(_has_captures) {}

    void Generate(InitsManager* im, std::vector<void*>& inits_vec, int offset) const override {
        auto l = make_intrusive<T>(name);
        auto& ft = im->Types(func_type);
        register_lambda__CPP(l, h, name, ft, has_captures);
    }

protected:
    const char* name;
    int func_type;
    p_hash_type h;
    bool has_captures;
};

// Constructs at run-time a mapping between abstract record field offsets used
// when compiling a set of scripts to their concrete offsets (which might differ
// from those during compilation due to loading of other scripts that extend
// various records).
class CPP_FieldMapping {
public:
    CPP_FieldMapping(int _rec, std::string _field_name, int _field_type, int _field_attrs)
        : rec(_rec), field_name(std::move(_field_name)), field_type(_field_type), field_attrs(_field_attrs) {}

    int ComputeOffset(InitsManager* im) const;

private:
    int rec;                // index to retrieve the record's type
    std::string field_name; // which field this offset pertains to
    int field_type;         // the field's type, in case we have to construct it
    int field_attrs;        // the same for the field's attributes
};

// Constructs at run-time a mapping between abstract enum values used when
// compiling a set of scripts to their concrete values (which might differ
// from those during compilation due to loading of other scripts that extend
// the enum).
class CPP_EnumMapping {
public:
    CPP_EnumMapping(int _e_type, std::string _e_name) : e_type(_e_type), e_name(std::move(_e_name)) {}

    int ComputeOffset(InitsManager* im) const;

private:
    int e_type;         // index to EnumType
    std::string e_name; // which enum constant for that type
};

// Looks up a BiF of the given name, making it available to compiled
// code via a C++ global.
class CPP_LookupBiF {
public:
    CPP_LookupBiF(zeek::Func*& _bif_func, std::string _bif_name)
        : bif_func(_bif_func), bif_name(std::move(_bif_name)) {}

    void ResolveBiF() const {
        // We allow a BiF to be resolved multiple times.  This is to
        // support plugins that might load BiFs that aren't initially
        // available, and also global initializations that might
        // require (other) BiFs that *are* initially available.
        if ( ! bif_func )
            bif_func = lookup_bif__CPP(bif_name.c_str());
    }

protected:
    zeek::Func*& bif_func; // where to store the pointer to the BiF
    std::string bif_name;  // the BiF's name
};

// Information needed to register a compiled function body (which makes it
// available to substitute for the body's AST).  The compiler generates
// code that loops over a vector of these to perform the registrations.
struct CPP_RegisterBody {
    CPP_RegisterBody(std::string _func_name, void* _func, int _type_signature, int _priority, p_hash_type _h,
                     const char* _filename, int _line_num, std::vector<std::string> _events)
        : func_name(std::move(_func_name)),
          func(_func),
          type_signature(_type_signature),
          priority(_priority),
          h(_h),
          filename(_filename),
          line_num(_line_num),
          events(std::move(_events)) {}

    std::string func_name; // name of the function
    void* func;            // pointer to C++
    int type_signature;
    int priority;
    p_hash_type h;
    const char* filename;
    int line_num;
    std::vector<std::string> events;
};

} // namespace zeek::detail