context.hpp

#pragma once

#include <string>
#include <unordered_map>
#include <stack>
#include <cassert>
#include <vector>
#include <set>


std::vector<std::string_view> tokenize_type(const std::string_view& type_str) {
        std::vector<std::string_view> tokens;
        if (type_str.empty()) {
            return tokens;
        }

        // tokenize by spaces NO REGEX
        size_t cursor = 0;
        size_t token_start = 0;

        while (cursor < type_str.length()) {
            if (type_str[cursor] == ' ') {

                assert(token_start < cursor);
                tokens.push_back(type_str.substr(token_start, cursor - token_start));
                token_start = cursor + 1;
            
            } else if (type_str[cursor] == '&' || type_str[cursor] == '*') {

                if (token_start < cursor) {
                    tokens.push_back(type_str.substr(token_start, cursor - token_start)); // push previous token
                }

                tokens.push_back(type_str.substr(cursor, 1)); // push & or * as separate token
                token_start = cursor + 1;

            } else if (type_str[cursor] == '<') {
                // skip untill matching > get inner string and recursively tokenize
                // push previous token
                
                if (token_start < cursor) {
                    tokens.push_back(type_str.substr(token_start, cursor - token_start));
                }



                size_t angle_bracket_count = 1;
                size_t inner_start = cursor + 1;
                size_t inner_cursor = inner_start;
                while (inner_cursor < type_str.length() && angle_bracket_count > 0) {
                    if (type_str[inner_cursor] == '<') {
                        angle_bracket_count++;
                    } else if (type_str[inner_cursor] == '>') {
                        angle_bracket_count--;
                    }
                    inner_cursor++;
                }

                assert(angle_bracket_count == 0);
                // std::string_view inner_string = std::string_view(type_str.c_str() + inner_start, inner_cursor - inner_start - 1);

                auto inner_tokens = tokenize_type(type_str.substr(inner_start, inner_cursor - inner_start - 1));
                
                // insert < token
                tokens.push_back(type_str.substr(cursor, 1));
                
                tokens.insert(tokens.end(), inner_tokens.begin(), inner_tokens.end());

                tokens.push_back(type_str.substr(inner_cursor - 1, 1));

                token_start = inner_cursor; // move token start after >
                cursor = inner_cursor - 1; // move cursor to >
            }
            
            cursor++;

        }
        
        if (token_start < cursor) {
            tokens.push_back(type_str.substr(token_start, cursor - token_start));
        }

        return tokens;
}


std::set<std::string> built_in_types = {
    "void", "bool", "char", "signed char", "unsigned char",
    "short", "unsigned short", "int", "unsigned int",
    "long", "unsigned long", "long long", "unsigned long long",
    "float", "double", "long double",
    "wchar_t", "char16_t", "char32_t",
    "int8_t", "uint8_t", "int16_t", "uint16_t", "int32_t", "uint32_t",
    "int64_t", "uint64_t", "float16", "bfloat16"
};

std::set<std::string> type_modifiers = {
    "const", "volatile", "mutable", "static", "extern", "register",
    "signed", "unsigned", "long", "short", "&", "*", "&", "<", ">"
};

struct namespace_t {
    
    std::string name;
    std::unordered_map<std::string, std::string> aliases;
    std::unordered_map<std::string, std::string_view> definitions;

    namespace_t(std::string n) : name(std::move(n)) {}

    // make it movable not copyable
    namespace_t(const namespace_t&) = delete;
    namespace_t& operator=(const namespace_t&) = delete;
    namespace_t(namespace_t&&) = default;
    namespace_t& operator=(namespace_t&&) = default;


    std::string resolve_alias(const std::string& alias_name) const {
        if (aliases.find(alias_name) != aliases.end()) {
            return resolve_alias(aliases.at(alias_name));
        }
        
        return alias_name;
    }

};

struct context_t {

    // define context_t as movable not copyable
    context_t(const context_t&) = delete;
    context_t& operator=(const context_t&) = delete;
    context_t(context_t&&) = default;
    context_t& operator=(context_t&&) = default;

    std::string source_code;
    std::unordered_map<std::string, std::string> normalization_cache;
    std::unordered_map<std::string, namespace_t> namespaces;

    context_t(std::string __src = "") {

        this->source_code = std::move(__src);

        // std::cerr << src << "\n";
        // src = "struct A { const char* bug = \" This will led to a bug }}}}\\\"}}} \"; }; \n";
        // std::cout << "//" << src << "\n";

        // open file test.cpp and read content
        // std::ifstream file("test.cpp");
        // assert(file.is_open());
        // if (file) {
        //     std::stringstream buffer;
        //     buffer << file.rdbuf();
        //     src = buffer.str();
        // }

        const std::string &src = this->source_code;

        struct scope_t{
            std::string namespace_name;
            size_t par_count;
            scope_t(const std::string& ns) : namespace_name(ns), par_count(0) {}
        };

        namespaces.emplace("global", namespace_t("global"));
        

        std::stack<scope_t> scope_stack;
        scope_stack.push(scope_t("global"));
        scope_stack.top().par_count = 1;

        size_t cursor = 0;
        auto skip_untill_next_token = [](const std::string& s, size_t& cursor) {
            size_t prev_cursor = cursor;
            // while (cursor < s.length() && s[cursor] != ' ' && s[cursor] != '\n' && s[cursor] != '\t') {
            while (cursor < s.length() && (s[cursor] == ' ' || s[cursor] == '\n' || s[cursor] == '\t')) {
                cursor++;
            }

            // std::cout << s.substr(prev_cursor, cursor - prev_cursor);

            return s[cursor] != ' ' && s[cursor] != '\n' && s[cursor] != '\t';
        };

        auto next_token = [](const std::string& s, size_t& cursor) -> std::string_view {
            if (s[cursor] == ' ' || s[cursor] == '\n' || s[cursor] == '\t') {
                return std::string_view();
            }   


            size_t start = cursor;
            while (cursor < s.length() && s[cursor] != ' ' && s[cursor] != '\n' && s[cursor] != '\t') {
                if (s[cursor] == ';' || s[cursor] == '{' || s[cursor] == '}' || s[cursor] == '(' || s[cursor] == ')' || s[cursor] == '<' || s[cursor] == '>') {
                    if (cursor == start) {
                        cursor++;
                    }
                    break;
                } else if (s[cursor] == '"' || s[cursor] == '\'') {

                    const char quote_char = s[cursor];
                    bool escaping = false;
                    while (cursor < s.length() - 1) {
                        cursor++;
                        if (s[cursor] == quote_char && !escaping) {
                            break;
                        } else if (s[cursor] == '\\' && !escaping) {
                            escaping = true;
                        } else {
                            escaping = false;
                        }
                    }
                }
                cursor++;
            }

            // print consumed char to stdout
            // std::cout << s.substr(start, cursor - start);

            return std::string_view(s.c_str() + start, cursor - start);
        };


        while (cursor < src.length()) {
            assert(scope_stack.size() >= 1);
            // std::cout << std::flush;
            size_t prev_cursor = cursor;

            // std::cout << "Cursor at: " << cursor << "\n";
            skip_untill_next_token(src, cursor);
            std::string_view token = next_token(src, cursor);
            if (token.empty()) continue;

            // std::cout << "Token: " << token << "\n";

            if (token.length() == 1) {
                switch (token[0])
                {

                    case '{':
                        // std::cout << "/*" << scope_stack.top().par_count << "*/";
                        scope_stack.top().par_count++;
                        break;
                    case '}':
                        // std::cout << "/*" << scope_stack.top().par_count  - 1 << "*/";
                        // std::cout << "} namespace: " << scope_stack.top().ns->name << " par_count=" << scope_stack.top().par_count - 1 << std::endl;
                        assert(scope_stack.size() >= 1);
                        assert(scope_stack.top().par_count >= 0);
                        scope_stack.top().par_count--;
                        if (scope_stack.top().par_count == 0) {
                            // std::cout << "/*-" << scope_stack.top().ns->name << "*/" << std::flush;
                            scope_stack.pop();
                        }
                        break;
                    default:
                        break;

                }

                continue;
            }

            // std::cout << "Token: " << token << "\n";

            if (token == "namespace") {
                assert(skip_untill_next_token(src, cursor));
                std::string_view ns_name = next_token(src, cursor);
                std::string full_ns_name = scope_stack.top().namespace_name == "global" ? std::string(ns_name) : scope_stack.top().namespace_name + "::" + std::string(ns_name);

                if (namespaces.find(full_ns_name) == namespaces.end()) {
                    namespaces.emplace(full_ns_name, namespace_t(full_ns_name));
                }

                scope_stack.push(scope_t(full_ns_name));

                // std::cout << "/*" << full_ns_name << "*/" << std::flush;

                // std::cerr << "Entering namespace: " << full_ns_name << std::endl;
            }

            else if (token == "template") {
                // skip template parameters
                size_t start_curs = cursor;
                size_t angle_bracket_count = 0;

                assert(skip_untill_next_token(src, cursor));
                std::string_view first_token = next_token(src, cursor);
                if (first_token == "<") 
                    angle_bracket_count++;
                else 
                    cursor = start_curs; // rewind (THIS IS NEEDED FOR extern template)
                
                while (angle_bracket_count != 0) {
                    assert(skip_untill_next_token(src, cursor));
                    std::string_view ttoken = next_token(src, cursor);
                    if (ttoken == "<") {
                        angle_bracket_count++;
                    } else if (ttoken == ">") {
                        angle_bracket_count--;
                    }
                }
            
                // print string (start_curs, cursor)
                std::string template_params = src.substr(start_curs, cursor - start_curs);
            }

            else if (token == "using") {
                
            //     // 1 Declaration
            //     // 1.1 Namespace member: es. using namespace_name::member_name
            //     // 1.2 Base class member: es. using BaseClass::member_name
            //     // 2 Directive
            //     // 2.1 Namespace directive: es. using namespace namespace_name
            //     // 3 Alias: es. using alias_name = actual_name
            //     // 4 inside enum class: using BaseEmum::field;

            //     // Templates could cause PROBLEMS here...
                std::string_view next_tok;
                assert(skip_untill_next_token(src, cursor));
                next_tok = next_token(src, cursor);
                if (next_tok == "namespace") {
                    continue;
                }

                std::string using_name = std::string(next_tok);
         
                assert(skip_untill_next_token(src, cursor));
                next_tok = next_token(src, cursor);    

                // if starts with :: skip it
                if (using_name.substr(0, 2) == "::") {
                    // check if there is another ::
                    size_t pos = using_name.find("::", 2);
                    if (pos != std::string::npos) {
                        using_name = using_name.substr(2);
                    } else {
                        continue;   
                    }
                }

        
                if (next_tok == "=") {
                    // Alias
                    assert(skip_untill_next_token(src, cursor));
                    std::string_view actual_name = next_token(src, cursor);
                    namespaces.at(scope_stack.top().namespace_name).aliases.emplace(using_name, std::string(actual_name));

                    // std::cerr << "Alias: " << using_name << " = " << actual_name << " in namespace " << scope_stack.top().namespace_name << "\n";
                
                } else {
                    if (scope_stack.top().par_count > 1) {
                        continue;
                    } else {
                        // Declaration namespace member!! (since not in other scopes)
                        std::string member_name = using_name.substr(using_name.find_last_of("::") + 1);
                        namespaces.at(scope_stack.top().namespace_name).aliases.emplace(member_name, using_name);
                        // std::cerr << "in namespace " << scope_stack.top().ns->name << " added alias " << member_name << " = " << using_name << "\n";

                        // std::cout << "Using namespace member: " << using_name << " in namespace " << scope_stack.top().ns->name << "\n";
                        // THIS IS NOT USEFULT FOR TYPE RESOLUTION
                    }
                    
                    // Declaration
                    // or using enum ...
                    // or using BaseEnum::field;
                    // std::cout << "Using: " << using_name << " in namespace " << scope_stack.top().ns->name << "\n";
                }
            }

            else if (token == "class" || token == "struct" || token == "enum" || token == "typedef") {

                size_t initial_def_cursor = prev_cursor;
                std::string decl_type = std::string(token);
                std::string typedef_name = "";

                if (token == "typedef") {
                    assert(skip_untill_next_token(src, cursor));
                    std::string_view next_tok = next_token(src, cursor);

                    if (next_tok == "enum" || next_tok == "struct" || next_tok == "class") {
                        // std::cout << "[WARNING] Skipping typedef parsing for now: " << next_tok << "\n";
                        // find first parenthesis then go to latest one and get the name
                        size_t saved_cursor = cursor;
                        size_t paren_count = 0;
                        while (true) {
                            assert(skip_untill_next_token(src, cursor));
                            std::string_view ttoken = next_token(src, cursor);
                            if (ttoken == "{") {
                                paren_count++;
                            } else if (ttoken == "}") {
                                paren_count--;
                                if (paren_count == 0) {
                                    break;
                                }
                            }
                        }
                        
                        assert(skip_untill_next_token(src, cursor));
                        std::string_view ttoken = next_token(src, cursor);
                        typedef_name = std::string(ttoken);
                        decl_type = std::string(next_tok);

                        cursor = saved_cursor;

                    } else {
                        // std::cout << "[WARNING] Skipping typedef parsing for now: " << next_tok << "\n";
                        continue; // skip other typedefs
                    }

                }  
                
                if (decl_type == "enum") {
                    size_t saved_cursor = cursor;
                    assert(skip_untill_next_token(src, cursor));
                    std::string_view ntoken = next_token(src, cursor);
                    if (ntoken == "class") {
                        decl_type = "enum class";
                    } else {
                        cursor = saved_cursor; // rewind
                    }
                }

                size_t saved_cursor = cursor;
               
                // since a class def could also be class __attribute__(...) Name {....}
                // or class Name : public Base { ... }
                // or class __attribute__(...) Name : public Base { ... }

                // to find the real name we must skip until we find : or { or ; and the token before is the name

                std::string_view class_name;
                while (true) {
                    size_t saved_cursor = cursor;
                    assert(skip_untill_next_token(src, cursor));
                    std::string_view ttoken = next_token(src, cursor);
                    if (ttoken == "{" || ttoken == ":" || ttoken == ";") {
                        cursor = saved_cursor; // rewind
                        break;
                    } else {
                        class_name = ttoken;
                    }
                }


                // assert(skip_untill_next_token(src, cursor));
                // std::string_view class_name = next_token(src, cursor);


                // if (class_name == "{") {
                //     cursor = saved_cursor; // rewind
                //     if (class_name.empty() && typedef_name.empty()) 
                //         continue; // anonymous class/struct/enum
                // }   

                // // extract string_view of class definition, find end index wich is the last }
                // // find first opening brace
                // // std::cerr << "Extracting class definition for: " << decl_type << " " << class_name << " | typedef " << typedef_name << "\n";


                // // print source code around +50 chars and -50 chars
                // if (decl_type == "class") {
                //     std::cerr << "\n\n-------------------- CODE CONTEXT --------------------\n";


                //     size_t context_start = initial_def_cursor >= 100 ? initial_def_cursor - 100 : 0;
                //     size_t context_end = (cursor + 100) < src.length() ? (cursor + 100) : src.length() - 1;
                //     std::cerr << "...\n" << src.substr(context_start, context_end - context_start) << "\n...\n";

                //     std::cerr << "-------------------- CODE CONTEXT --------------------\n";

                // }

                size_t brace_count = 0;
                while (cursor < src.length()) {
                    assert(skip_untill_next_token(src, cursor));
                    std::string_view ttoken = next_token(src, cursor);
                    // std::cerr << "Token in class body: " << ttoken << "\n";
                    if (ttoken == "{") {
                        brace_count++;
                    } else if (ttoken == "}") {
                        assert(brace_count > 0);
                        brace_count--;
                        if (brace_count == 0) {
                            break;
                        }
                    } else if (ttoken == ";") {
                        // forward declaration
                        if (brace_count == 0) {
                            break;
                        }
                    }
                }


                std::string_view class_def = std::string_view(src).substr(initial_def_cursor, cursor - initial_def_cursor + 1);
                // std::string_view class_def = src.substr(initial_def_cursor, cursor - initial_def_cursor + 1);

                // if (decl_type == "class") {
                //     std::cerr << "Extracting class definition for: " << decl_type << " " << class_name << " | typedef " << typedef_name << "\n";
                //     std::cerr << class_def << "\n";
                // }

                std::cerr << std::flush;

                // store definition
                if (!typedef_name.empty()) {
                    namespaces.at(scope_stack.top().namespace_name).definitions.emplace(typedef_name, class_def);
                } 

                if (!class_name.empty()) {
                    namespaces.at(scope_stack.top().namespace_name).definitions.emplace(std::string(class_name), class_def);
                }
            }

        }

        // print tvm::runtime::vm::VirtualMachine

    }

    // print scopes 
    void print_all_scopes() {
        std::cerr << "Namespaces in context:\n";
        for (const auto& [ns_name, ns] : namespaces) {
            std::cerr << ns_name << ", ";
        }
        std::cerr << "\n";
    }


    const std::string_view get_source_code_of(const std::string& name) const {

        std::string ns = "global";
        if (name.find("::") != std::string::npos) {
            size_t pos = name.rfind("::");
            ns = name.substr(0, pos);
        }

        std::string simple_name = name.substr(name.find_last_of("::") + 1);
        if (namespaces.find(ns) == namespaces.end()) {
            return std::string_view();
        }

        if (namespaces.at(ns).definitions.find(simple_name) == namespaces.at(ns).definitions.end()) {
            return std::string_view();
        }

        return namespaces.at(ns).definitions.at(simple_name);
    }


    std::string get_abs_scope(const std::string& current_scope, const std::string& name) const {
        // std::cerr << "Resolving absolute scope for name: " << name << " from current scope: " << current_scope << "\n";
        if (name.length() < 2) 
            return current_scope;

        // if name starts with :: it is global
        if (name.substr(0, 2) == "::") 
            return "global";

        // look if :: present in name
        size_t dd_pos = name.rfind("::");
        if (dd_pos == std::string::npos) 
            return current_scope;

        std::string type_namespace = name.substr(0, dd_pos);
        // std::cerr << "Type namespace extracted: " << type_namespace << "\n";

        // // debug print all namespaces
        // std::cerr << "Available namespaces: " << namespaces.size() << "\n";
        // for (const auto& [ns_name, ns] : namespaces) {
        //     std::cerr << ns_name << ", ";
        // }
        // std::cerr << "\n" << std::flush;

        if (namespaces.find(type_namespace) != namespaces.end()) {
            // std::cerr << "Type namespace found directly: " << type_namespace << "\n";
            return type_namespace;
        }

        // std::vector<std::string_view> ns_parts;
        // size_t start = 0;
        // size_t pos = 0;
        // while ((pos = type_namespace.find("::", start)) != std::string::npos) {
        //     ns_parts.push_back(std::string_view(type_namespace).substr(start, pos - start));
        //     start = pos + 2;
        // }
        // ns_parts.push_back(std::string_view(type_namespace).substr(start));

        std::string search_scope = current_scope;
        while (search_scope != "") {
            
            // is namespace a child nemespace of search_scope?
            std::string full_ns = search_scope == "global" ? type_namespace : search_scope + "::" + type_namespace;
            if (namespaces.find(full_ns) != namespaces.end()) {
                return full_ns;
            }

            // go up in the namespace hierarchy split by ::
            size_t pos = search_scope.rfind("::");
            if (pos != std::string::npos) {
                search_scope = search_scope.substr(0, pos);
            } else if (search_scope != "global") {
                search_scope = "global";
            } else {
                search_scope = "";
            }
        }

        assert(false && "Namespace not found");
    }

    std::string get_abs_def(const std::string& scope, const std::string& name) const {
        assert(namespaces.find(scope) != namespaces.end() && "Scope not found in context");
        std::string abs_scope = get_abs_scope(scope, name);
        std::string simple_name = name.substr(name.find_last_of("::") + 1);

        const namespace_t& ns = namespaces.at(abs_scope);

        if (ns.definitions.find(simple_name) != ns.definitions.end()) {
            return abs_scope != "global" ? abs_scope + "::" + simple_name : simple_name;
        }

        std::string resolved_alias = ns.resolve_alias(simple_name);
   
        
        if (resolved_alias != simple_name) {
            return get_abs_def(scope, resolved_alias);
        } else if (scope == "global") {
            return "";
        } else {
            // try go up in the namespace hierarchy split by :: if no parent nemespace look in global
            size_t pos = scope.rfind("::");
            if (pos != std::string::npos) {
                std::string parent_scope = scope.substr(0, pos);
                return get_abs_def(parent_scope, resolved_alias);
            } else {
                return get_abs_def("global", resolved_alias);
            }
        }
     
    }

    std::string normalize_type(const std::string& from_namespace, const std::string& type_str) const {
        auto tokens = tokenize_type(type_str);

        // std::string result = "OriginalType: " + type_str + "\nNormalizedType:";
        std::string result;
        for (size_t i = 0; i < tokens.size(); ++i) {
            // skip modifiers
            if (type_modifiers.find(std::string(tokens[i])) != type_modifiers.end()) {
                // result += "\033[1m" + std::string(tokens[i]) + "\033[0m ";
                result += std::string(tokens[i]) + " ";
                continue;
            }

            // check built-in types
            if (built_in_types.find(std::string(tokens[i])) != built_in_types.end()) {
                // result += "\033[4m" + std::string(tokens[i]) + "\033[0m ";
                result += std::string(tokens[i]) + " ";
                continue;
            }

            // Decode the type from context
            std::string normalized = this->get_abs_def(from_namespace, std::string(tokens[i]));
            if (!normalized.empty()) {
                result += normalized + " ";
                continue;
            }

            return "";
        }

        return result;
    }

};