Programming in Lua. Exercise 29.4

Lua.Exercise29.4.cpp

// Exercise 29.4:
// Based on the example for boolean arrays,
// implement a small C library for integer arrays.

#define LUA_COMPAT_APIINTCASTS

#include <cmath>
#include <cstdint>
#include <exception>
#include <iomanip>
#include <iostream>
#include <limits>
#include <string>
#include <vector>
#include <stack>
#include <unordered_set>

#include "lua.hpp"

namespace {

struct IntArray {
    size_t size;
    int32_t values[1];
};

#define checkintarray(L, i) \
    (IntArray*) luaL_checkudata(L, i, "LuaType.intarray")

bool AbortWithMessage(const char* message) {
    std::cout << message << std::endl;
    std::abort();
}

bool LoadFile(lua_State* state, const std::string& filename) {
    return !(luaL_loadfile(state, filename.c_str()) || lua_pcall(state, 0, 0, 0));
}

IntArray* intarray_new(lua_State* state, size_t size) {
    size_t nbytes = sizeof(IntArray) + (size - 1) * sizeof(int32_t);
    IntArray* array = (IntArray*) lua_newuserdata(state, nbytes);

    array->size = size;
    for (size_t i = 0; i < size; i++) {
        array->values[i] = 0;
    }

    luaL_getmetatable(state, "LuaType.intarray");
    lua_setmetatable(state, -2);

    return array;
}

int lua_intarray_new(lua_State* state) {
    int size = luaL_checkint(state, 1);
    luaL_argcheck(state, size >= 1, 1, "invalid size");
    intarray_new(state, size);
    return 1;
}

int lua_intarray_setitem(lua_State* state) {
    IntArray* a = checkintarray(state, 1);
    int index = luaL_checkint(state, 2) - 1;
    luaL_argcheck(state, 0 <= index && index < a->size, 2, "index out of range.");
    a->values[index] = static_cast<int32_t>(luaL_checkint(state, 3));
    return 0;
}

int lua_intarray_getitem(lua_State* state) {
    IntArray* a = checkintarray(state, 1);
    int index = luaL_checkint(state, 2) - 1;
    luaL_argcheck(state, 0 <= index && index < a->size, 2, "index out of range.");
    lua_pushinteger(state, a->values[index]);
    return 1;
}

int lua_intarray_getsize(lua_State *L) {
    IntArray* a = checkintarray(L, 1);
    lua_pushinteger(L, a->size);
    return 1;
}

int lua_intarray_tostring(lua_State* state) {
    IntArray* a = checkintarray(state, 1);

    luaL_Buffer buffer;
    char* cbuffer = luaL_buffinitsize(state, &buffer, /* length= */ a->size * 11 + 2);
    size_t cbuffer_index = 0;

    cbuffer[cbuffer_index] = '[';
    cbuffer_index += 1;

    for (size_t i = 0; i < a->size; i++) {
        int32_t value = a->values[i];

        if (value == 0) {
            cbuffer[cbuffer_index] = '0';
            cbuffer_index += 1;
        } else {
            // value <> 0
            if (value < 0) {
                cbuffer[cbuffer_index] = '-';
                cbuffer_index += 1;
                value *= -1;
            }

            std::stack<int32_t> numbers;
            while (value > 0) {
                numbers.push(value % 10);
                value /= 10;
            }

            while (!numbers.empty()) {
                cbuffer[cbuffer_index] = static_cast<char>(numbers.top() + '0');
                numbers.pop();
                cbuffer_index += 1;
            }
        }

        if (i + 1 < a->size) {
            cbuffer[cbuffer_index + 0] = ',';
            cbuffer[cbuffer_index + 1] = ' ';
            cbuffer_index += 2;
        }
    }

    cbuffer[cbuffer_index] = ']';
    cbuffer_index += 1;

    luaL_pushresultsize(&buffer, cbuffer_index);
    return 1;
}

int lua_intarray_union(lua_State* state) {
    IntArray* a = checkintarray(state, 1);
    IntArray* b = checkintarray(state, 2);

    std::unordered_set<int32_t> lookup;
    for (size_t i = 0; i < a->size; i++) {
        lookup.insert(a->values[i]);
    }

    size_t u_size = 0;
    std::unordered_set<int32_t> used;
    for (size_t i = 0; i < b->size; i++) {
        if (lookup.find(b->values[i]) != lookup.end() &&
            used.find(b->values[i]) == used.end()) {
            used.insert(b->values[i]);
            u_size += 1;
        }
    }

    int32_t u_index = 0;
    IntArray* u = intarray_new(state, u_size);

    used.clear();
    for (size_t i = 0; i < b->size; i++) {
        if (lookup.find(b->values[i]) != lookup.end() &&
            used.find(b->values[i]) == used.end()) {
            used.insert(b->values[i]);
            u->values[u_index] = b->values[i];
            u_index += 1;
        }
    }

    return 1;
}

int lua_intarray_intersection(lua_State* state) {
    IntArray* a = checkintarray(state, 1);
    IntArray* b = checkintarray(state, 2);

    std::unordered_set<int32_t> lookup;
    for (size_t i = 0; i < a->size; i++) {
        lookup.insert(a->values[i]);
    }

    for (size_t i = 0; i < b->size; i++) {
        lookup.insert(b->values[i]);
    }

    int32_t u_index = 0;
    IntArray* u = intarray_new(state, lookup.size());

    for (int32_t value: lookup) {
        u->values[u_index] = value;
        u_index += 1;
    }

    return 1;
}

const struct luaL_Reg intarrayLib_f[] = {
    {"new", lua_intarray_new},
    {NULL, NULL}
};

const struct luaL_Reg intarrayLib_m[] = {
    {"__tostring", lua_intarray_tostring},
    {"set", lua_intarray_setitem},
    {"get", lua_intarray_getitem},
    {"size", lua_intarray_getsize},
    {"union", lua_intarray_union},
    {"intersection", lua_intarray_intersection},
    {NULL, NULL}
};

LUALIB_API int luaopen_intarrayLib(lua_State* state) {
    luaL_newmetatable(state, "LuaType.intarray");

    lua_pushvalue(state, -1);
    lua_setfield(state, -2, "__index");

    luaL_setfuncs(state, intarrayLib_m, 0);

    luaL_newlib(state, intarrayLib_f);
    lua_setglobal(state, "intarray");

    return 0;
}

}  // namespace

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cout << "Please, specify lua script file to run." << std::endl;
        return 0;
    }

    lua_State* L = luaL_newstate();
    luaL_openlibs(L);
    luaopen_intarrayLib(L);

    if (!LoadFile(L, std::string(argv[1]))) {
        AbortWithMessage(lua_tostring(L, -1));
        return -1;
    }

    return 0;
}

main.lua

a = intarray.new(10)
b = intarray.new(5)

a:set(1, -7)
a:set(5, 22)
a:set(7, 31)
a:set(8, 108)

print("a:", a)

b:set(1, 55)
b:set(2, -7)
b:set(5, 22)

print("b:", b)

print("intersection:", a:intersection(b))
print("union:", a:union(b))

output.txt

a:	[-7, 0, 0, 0, 22, 0, 31, 108, 0, 0]
b:	[55, -7, 0, 0, 22]
intersection:	[-7, 108, 31, 55, 22, 0]
union:	[-7, 0, 22]