utils.cc

/*
 *  MoMEMta: a modular implementation of the Matrix Element Method
 *  Copyright (C) 2016  Universite catholique de Louvain (UCL), Belgium
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <lua/utils.h>

#include <cassert>
#include <regex>

#include <momemta/InputTag.h>
#include <momemta/ILuaCallback.h>
#include <momemta/Logging.h>
#include <momemta/ModuleRegistry.h>
#include <momemta/ParameterSet.h>
#include <momemta/Utils.h>

#include <LibraryManager.h>
#include <lua/ParameterSetParser.h>
#include <lua/bindings/Path.h>
#include <lua/bindings/Types.h>

// Defined by `embedLua.py` at build-time
extern void execute_embed_lua_code(lua_State*);

namespace lua {

    Lazy::Lazy(lua_State* L) {
        this->L = L;
    }

    LazyFunction::LazyFunction(lua_State* L, int index): Lazy(L) {
        auto absolute_index = get_index(L, index);

        // Duplicate the function on the top of the stack. This ensure the stack size won't change
        lua_pushvalue(L, absolute_index);

        // Pop the anonymous function from the stack, and store it in the global lua registry
        ref_index = luaL_ref(L, LUA_REGISTRYINDEX);
    }

    momemta::any LazyFunction::operator() () const {

        LOG(trace) << "[LazyFunction::operator()] >> stack size = " << lua_gettop(L);

        // Pop the anonymous function from the registry, and push it on the top of the stack
        lua_rawgeti(L, LUA_REGISTRYINDEX, ref_index);

        // Call the function. The function removed from the stack, and the return value is pushed on the top of the stack
        auto result = lua_pcall(L, 0, 1, 0);
        if (result != LUA_OK) {
            std::string error = lua_tostring(L, -1);
            LOG(fatal) << "Fail to call lua anonymous function. Return value is " << result << ". Error message: " << error;
        }

        momemta::any value;
        bool lazy = false;
        std::tie(value, lazy) = to_any(L, -1);
        assert(!lazy);

        lua_pop(L, 1);

        LOG(trace) << "[LazyFunction::operator()] << stack size = " << lua_gettop(L);

        return value;
    }

    Type type(lua_State* L, int index) {
        int t = lua_type(L, index);

        switch (t) {
            case LUA_TBOOLEAN:
                return BOOLEAN;
                break;

            case LUA_TSTRING: {
                std::string value = lua_tostring(L, index);
                if (InputTag::isInputTag(value))
                    return INPUT_TAG;
                else
                    return STRING;
            } break;

            case LUA_TNUMBER:
                if (lua_isinteger(L, index))
                    return INTEGER;
                else
                    return REAL;
                break;

            case LUA_TTABLE: {
                // We only support ParameterSet for table
                if (lua_is_array(L, index) == -1) {
                    return PARAMETER_SET;
                }

            } break;
        }

        return NOT_SUPPORTED;
    }

    size_t get_index(lua_State* L, int index) {
        return (index < 0) ? lua_gettop(L) + index + 1 : index;
    }

    int lua_is_array(lua_State* L, int index) {

        LOG(trace) << "[lua_is_array] >> stack size = " << lua_gettop(L);

        size_t table_index = get_index(L, index);

        if (lua_type(L, table_index) != LUA_TTABLE)
            return -1;

        size_t size = 0;

        lua_pushnil(L);
        while (lua_next(L, table_index) != 0) {
            if (lua_type(L, -2) != LUA_TNUMBER) {
                lua_pop(L, 2);
                LOG(trace) << "[lua_is_array] << stack size = " << lua_gettop(L);
                return -1;
            }

            size++;

            lua_pop(L, 1);
        }

        LOG(trace) << "[lua_is_array] << stack size = " << lua_gettop(L);
        return size;
    }

    Type lua_array_unique_type(lua_State* L, int index) {

        size_t absolute_index = get_index(L, index);

        if (lua_type(L, absolute_index) != LUA_TTABLE)
            return NOT_SUPPORTED;

        Type result = NOT_SUPPORTED;

        lua_pushnil(L);
        while (lua_next(L, absolute_index) != 0) {
            if (result == NOT_SUPPORTED) {
                result = type(L, -1);
            } else {
                Type entry_type = type(L, -1);

                if ((result == INTEGER) && (entry_type == REAL))
                    result = REAL;
                else if ((result == REAL) && (entry_type == INTEGER))
                    result = REAL;
                else if (result != entry_type) {
                    lua_pop(L, 2);
                    result = NOT_SUPPORTED;
                    break;
                }
            }

            lua_pop(L, 1);
        }

        return result;
    }

    std::pair<momemta::any, bool> to_any(lua_State* L, int index) {

        LOG(trace) << "[to_any] >> stack size = " << lua_gettop(L);
        size_t absolute_index = get_index(L, index);

        momemta::any result;
        bool lazy = false;

        auto type = lua_type(L, absolute_index);
        switch (type) {
            case LUA_TNUMBER: {
                if (lua_isinteger(L, absolute_index)) {
                    int64_t number = lua_tointeger(L, absolute_index);
                    result = number;
                } else {
                    double number = lua_tonumber(L, absolute_index);
                    result = number;
                }
            } break;

            case LUA_TBOOLEAN: {
                bool value = lua_toboolean(L, absolute_index);
                result = value;
            } break;

            case LUA_TSTRING: {
                std::string value = lua_tostring(L, absolute_index);
                if (InputTag::isInputTag(value)) {
                    InputTag tag = InputTag::fromString(value);
                    result = tag;
                } else {
                    result = value;
                }
            } break;

            case LUA_TTABLE: {
                LOG(trace) << "[to_any::table] >> stack size = " << lua_gettop(L);
                if (lua::lua_is_array(L, absolute_index) > 0) {

                    Type type = NOT_SUPPORTED;

                    if ((type = lua::lua_array_unique_type(L, absolute_index)) != NOT_SUPPORTED) {
                        result = to_vector(L, absolute_index, type);
                    } else {
                        throw invalid_array_error("Various types stored into the array. This is not supported.");
                    }

                } else {
                    ParameterSet cfg;
                    ParameterSetParser::parse(cfg, L, absolute_index);
                    result = cfg;
                }
                LOG(trace) << "[to_any::table] << stack size = " << lua_gettop(L);
            } break;

            case LUA_TFUNCTION: {
                LOG(trace) << "[to_any::function] >> stack size = " << lua_gettop(L);

                result = LazyFunction(L, absolute_index);
                lazy = true;

                LOG(trace) << "[to_any::function] << stack size = " << lua_gettop(L);
            } break;

            case LUA_TUSERDATA: {
                result = get_custom_type_ptr(L, absolute_index);

            } break;

            default: {
                LOG(fatal) << "Unsupported lua type: " << lua_type(L, absolute_index);
                throw lua::invalid_configuration_file("");
            } break;
        }

        LOG(trace) << "[to_any] << final type = " << demangle(result.type().name());
        LOG(trace) << "[to_any] << stack size = " << lua_gettop(L);
        return {result, lazy};
    }

    void push_any(lua_State* L, const momemta::any& value) {
        LOG(trace) << "[push_any] >> stack size = " << lua_gettop(L);

        if (value.type() == typeid(int64_t)) {
            int64_t v = momemta::any_cast<int64_t>(value);
            lua_pushinteger(L, v);
        } else if (value.type() == typeid(double)) {
            double v = momemta::any_cast<double>(value);
            lua_pushnumber(L, v);
        } else if (value.type() == typeid(bool)) {
            bool v = momemta::any_cast<bool>(value);
            lua_pushboolean(L, v);
        } else if (value.type() == typeid(std::string)) {
            auto v = momemta::any_cast<std::string>(value);
            lua_pushstring(L, v.c_str());
        } else if (value.type() == typeid(InputTag)) {
            auto v = momemta::any_cast<InputTag>(value).toString();
            lua_pushstring(L, v.c_str());
        } else {
            LOG(fatal) << "Unsupported C++ value: " << demangle(value.type().name());
            throw lua::unsupported_type_error(demangle(value.type().name()));
        }

        LOG(trace) << "[push_any] << stack size = " << lua_gettop(L);
    }

    momemta::any to_vector(lua_State* L, int index, Type t) {
        switch (t) {
            case BOOLEAN:
                return to_vectorT<bool>(L, index);

            case STRING:
                return to_vectorT<std::string>(L, index);

            case INTEGER:
                return to_vectorT<int64_t>(L, index);

            case REAL:
                return to_vectorT<double>(L, index);

            case INPUT_TAG:
                return to_vectorT<InputTag>(L, index);

            case PARAMETER_SET:
                return to_vectorT<ParameterSet>(L, index);

            case NOT_SUPPORTED:
                break;
        }

        throw invalid_array_error("Unsupported array type");
    }

    template<> double special_any_cast(const momemta::any& value) {
        if (value.type() == typeid(int64_t))
            return static_cast<double>(momemta::any_cast<int64_t>(value));

        return momemta::any_cast<double>(value);
    }

    int module_table_newindex(lua_State* L) {
        lua_getmetatable(L, 1);
        lua_getfield(L, -1, "__type");

        const char* module_type = luaL_checkstring(L, -1);
        const char* module_name = luaL_checkstring(L, 2);

        // Remove field name from stack
        lua_pop(L, 1);

        // Validate module name
        // Format is: [a-zA-Z][a-zA-Z0-9_]*
        static std::regex name_regex("[a-zA-Z][a-zA-Z0-9_]*");

        if (! std::regex_match(module_name, name_regex)) {
            luaL_error(L, "invalid module name '%s': valid format is [a-zA-Z][a-zA-Z0-9_]*", module_name);
        }

        lua_getfield(L, -1, "__ptr");
        void* cfg_ptr = lua_touserdata(L, -1);
        ILuaCallback* callback = static_cast<ILuaCallback*>(cfg_ptr);

        callback->onModuleDeclared(module_type, module_name);

        // Remove metatable from the stack
        lua_pop(L, 2);

        // Add "@name" and "@type" fields to the module's parameters

        // Push the key and then the value
        lua_pushstring(L, "@name");
        lua_pushstring(L, module_name);
        lua_rawset(L, -3);

        lua_pushstring(L, "@type");
        lua_pushstring(L, module_type);
        lua_rawset(L, -3);

        // And actually set the value to the table
        lua_rawset(L, 1);

        return 0;
    }

    void register_modules(lua_State* L, void* ptr) {
        momemta::ModuleList modules;
        momemta::ModuleRegistry::get().exportList(true, modules);

        for (const auto& module: modules) {
            const char* module_name = module.name.c_str();

            int type = lua_getglobal(L, module_name);
            lua_pop(L, 1);
            if (type != LUA_TNIL) {
                // Global already exists
                continue;
            }

            // Create a new empty table
            lua_newtable(L);

            std::string module_metatable = module.name + "_mt";

            // Create the associated metatable
            luaL_newmetatable(L, module_metatable.c_str());

            lua_pushstring(L, module_name);
            lua_setfield(L, -2, "__type");

            lua_pushlightuserdata(L, ptr);
            lua_setfield(L, -2, "__ptr");

            // Set the metatable '__newindex' function
            // This function is called when a assignment is made to the table
            // In our case, it's called when a new module is declared
            const luaL_Reg l[] = {
                {"__newindex", lua::module_table_newindex},
                {nullptr, nullptr}
            };
            luaL_setfuncs(L, l, 0);

            lua_setmetatable(L, -2);

            // And register it as a global variable
            lua_setglobal(L, module_name);

            LOG(trace) << "Registered new lua global variable '" << module_name << "'";
        }
    }

    int load_modules(lua_State* L) {
        int n = lua_gettop(L);
        if (n != 1) {
            luaL_error(L, "invalid number of arguments: 1 expected, got %d", n);
        }

        void* cfg_ptr = lua_touserdata(L, lua_upvalueindex(1));

        const char *path = luaL_checkstring(L, 1);
        LibraryManager::get().registerLibrary(path);

        register_modules(L, cfg_ptr);

        return 0;
    }

    int parameter(lua_State* L) {
        int n = lua_gettop(L);
        if (n != 1) {
            luaL_error(L, "invalid number of arguments: 1 expected, got %d", n);
        }

        std::string parameter_name = luaL_checkstring(L, 1);

        // Create an anonymous function return the value of the parameter
        // Assumes there's a global table named `parameters`

        std::string code = "return function() return parameters['" + parameter_name + "'] end";
        luaL_dostring(L, code.c_str());

        return 1;
    }

    int set_final_module(lua_State* L) {
        int n = lua_gettop(L);
        if (n == 0) {
            luaL_error(L, "invalid number of arguments: at least one expected, got 0");
        }

        void* cfg_ptr = lua_touserdata(L, lua_upvalueindex(1));
        ILuaCallback* callback = static_cast<ILuaCallback*>(cfg_ptr);

        for(size_t i = 1; i <= size_t(n); i++) {
            std::string input_tag = luaL_checkstring(L, i);
            if (!InputTag::isInputTag(input_tag)) {
                luaL_error(L, "'%s' is not a valid InputTag", input_tag.c_str());
            }
            callback->onIntegrandDeclared(InputTag::fromString(input_tag));
        }

        return 0;
    }

    int add_integration_dimension(lua_State* L) {
        int n = lua_gettop(L);
        if (n != 0) {
            luaL_error(L, "invalid number of arguments: 0 expected, got %d", n);
        }

        // Create input tag using current value of the index
        int64_t cuba_index = lua_tonumber(L, lua_upvalueindex(1));
        lua_pushnumber(L, cuba_index + 1);
        lua_replace(L, lua_upvalueindex(1));

        std::string index_tag = "cuba::ps_points/";
        index_tag += std::to_string(cuba_index);

        // Input tag is return value of the function
        push_any(L, index_tag);

        // Add an integration dimension in the configuration
        void* cfg_ptr = lua_touserdata(L, lua_upvalueindex(2));
        ILuaCallback* callback = static_cast<ILuaCallback*>(cfg_ptr);
        callback->addIntegrationDimension();

        return 1;
    }

    int declare_input(lua_State* L) {
        int n = lua_gettop(L);
        if (n != 1) {
            luaL_error(L, "invalid number of arguments: 1 expected, got %d", n);
        }

        std::string input_name = luaL_checkstring(L, 1);

        void* cfg_ptr = lua_touserdata(L, lua_upvalueindex(1));
        ILuaCallback* callback = static_cast<ILuaCallback*>(cfg_ptr);

        callback->onNewInputDeclared(input_name);

        return 0;
    }

    void setup_hooks(lua_State* L, void* ptr) {
        lua_pushlightuserdata(L, ptr);
        lua_pushcclosure(L, load_modules, 1);
        lua_setglobal(L, "load_modules");

        lua_pushlightuserdata(L, ptr);
        lua_pushcclosure(L, parameter, 1);
        lua_setglobal(L, "parameter");

        // Define the `add_dimension()` function in Lua and make it available in the global namespace.
        // See add_integration_dimension for more information.
        lua_pushnumber(L, 1);
        lua_pushlightuserdata(L, ptr);
        lua_pushcclosure(L, add_integration_dimension, 2);
        lua_setglobal(L, "add_dimension");

        // integrand() function
        lua_pushlightuserdata(L, ptr);
        lua_pushcclosure(L, set_final_module, 1);
        lua_setglobal(L, "integrand");

        // momemta_declare_input function
        lua_pushlightuserdata(L, ptr);
        lua_pushcclosure(L, declare_input, 1);
        lua_setglobal(L, "momemta_declare_input");

        // C++ -> lua bindings of some classes
        path_register(L, ptr);
    }

    std::shared_ptr<lua_State> init_runtime(ILuaCallback* callback) {

        std::shared_ptr<lua_State> L(luaL_newstate(), lua_close);
        luaL_openlibs(L.get());

        // Register hooks function, like `load_modules`
        lua::setup_hooks(L.get(), callback);

        // Register existing modules
        lua::register_modules(L.get(), callback);

        // Default functions
        // Function defined by `embedLua.py` at build-time
        execute_embed_lua_code(L.get());

        return L;
    }

    void inject_parameters(lua_State* L, const ParameterSet& parameters) {
        for (const auto& parameter: parameters.getNames()) {
            LOG(debug) << "Injecting parameter " << parameter;
            lua::push_any(L, parameters.rawGet(parameter));
            lua_setglobal(L, parameter.c_str());
        }
    }

    namespace debug {
        std::vector<std::string> dump_stack(lua_State *L) {
            std::vector<std::string> stack;
            for (int i = 1; i < lua_gettop(L) + 1; i++) {
                if (lua_isnumber(L, i)) {
                    stack.push_back("number : " + std::to_string(lua_tonumber(L, i)));
                } else if (lua_isstring(L, i)) {
                    stack.push_back(std::string("string : ") + std::string(luaL_checkstring(L, i)));
                } else if (lua_istable(L, i)) {
                    stack.push_back("table");
                } else if (lua_iscfunction(L, i)) {
                    stack.push_back("cfunction");
                } else if (lua_isfunction(L, i)) {
                    stack.push_back("function");
                } else if (lua_isboolean(L, i)) {
                    if (lua_toboolean(L, i) != 0)
                        stack.push_back("boolean: true");
                    else
                        stack.push_back("boolean: false");
                } else if (lua_isuserdata(L, i)) {
                    stack.push_back("userdata");
                } else if (lua_isnil(L, i)) {
                    stack.push_back("nil");
                } else if (lua_islightuserdata(L, i)) {
                    stack.push_back("lightuserdata");
                }
            }

            return stack;
        }

        void print_stack(lua_State* L) {
            auto stack = dump_stack(L);
            size_t index = 0;
            LOG(debug) << "Stack has " << stack.size() << " elements: ";
            for (const auto& e: stack) {
                LOG(debug) << "  #" << index++ << ": " << e;
            }
        }
    }
}