Programming in Lua. Exercise 29.2 & 29.3

Lua.Exercise29.2+3.cpp

// Exercise 29.2: 
// We can see a boolean array as a set of integers
// (the indices with true values in the array).
// Add to the implementation of boolean arrays
// functions to compute
// the union and intersection of two arrays.
// These functions should receive two arrays and return a new one,
// without modifying its parameters.

// Exercise 29.3: 
// Modify the implementation of the __tostring metamethod
// so that it shows the full contents of the array
// in an appropriate way.
// Use the buffer facility (the section called “String Manipulation”)
// to create the resulting string.

#define LUA_COMPAT_APIINTCASTS

#include <cmath>
#include <exception>
#include <iomanip>
#include <iostream>
#include <limits>
#include <string>
#include <vector>

#include "lua.hpp"

namespace {

#define BITS_PER_WORD (CHAR_BIT*sizeof(unsigned int))
#define I_WORD(i)     ((unsigned int)(i) / BITS_PER_WORD)
#define I_BIT(i)      (1 << ((unsigned int)(i) % BITS_PER_WORD))
#define N_WORDS(i)    ((unsigned int)(i) / BITS_PER_WORD + ((unsigned int)(i) % BITS_PER_WORD > 0 ? 1 : 0))

typedef struct NumArray {
    int size;
    unsigned int values[1];  /* variable part */
} NumArray;

#define checkarray(L, i) \
    (NumArray*) luaL_checkudata(L, i, "LuaBook.array")

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));
}

NumArray* array_new(lua_State* state, size_t size) {
    int i;
    size_t nbytes;
    NumArray* a;

    nbytes = sizeof(NumArray) + I_WORD(size - 1) * sizeof(unsigned int);
    a = (NumArray*) lua_newuserdata(state, nbytes);

    a->size = size;
    for (i = 0; i <= I_WORD(size - 1); i++) {
        a->values[i] = 0;
    }

    luaL_getmetatable(state, "LuaBook.array");
    lua_setmetatable(state, -2);

    return a;
}

int lua_array_new(lua_State* state) {
    int size = luaL_checkint(state, 1);
    luaL_argcheck(state, size >= 1, 1, "invalid size");
    NumArray* a = array_new(state, size);
    return 1;
}

unsigned int* getindex(lua_State* L, unsigned int* mask) {
    NumArray* a = checkarray(L, 1);
    int index = luaL_checkint(L, 2) - 1;
    
    luaL_argcheck(L, 0 <= index && index < a->size, 2,
                        "index out of range");
    
    /* return element address */
    *mask = I_BIT(index);
    return &a->values[I_WORD(index)];
}

int setarray(lua_State* L) {
    unsigned int mask;
    unsigned int* entry = getindex(L, &mask);
    
    if (!lua_isboolean(L, 3)) {
        luaL_error(L, "Expected boolean value.");
    }

    if (lua_toboolean(L, 3))
        *entry |= mask;
    else
        *entry &= ~mask;
    
    return 0;
}

int getarray(lua_State* L) {
    unsigned int mask;
    unsigned int* entry = getindex(L, &mask);
    lua_pushboolean(L, *entry & mask);
    return 1;
}

int getsize(lua_State *L) {
    NumArray* a = checkarray(L, 1);
    lua_pushinteger(L, a->size);
    return 1;
}

int array2string(lua_State* state) {
    NumArray* a = checkarray(state, 1);

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

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

    for (size_t i = 0; i < a->size; i++) {
        auto word = I_WORD(i);
        auto bit = I_BIT(i);
        auto value = (a->values[word] & bit);

        if (value) {
            cbuffer[cbuffer_index + 0] = 't';
            cbuffer[cbuffer_index + 1] = 'r';
            cbuffer[cbuffer_index + 2] = 'u';
            cbuffer[cbuffer_index + 3] = 'e';
            cbuffer_index += 4;
        } else {
            cbuffer[cbuffer_index + 0] = 'f';
            cbuffer[cbuffer_index + 1] = 'a';
            cbuffer[cbuffer_index + 2] = 'l';
            cbuffer[cbuffer_index + 3] = 's';
            cbuffer[cbuffer_index + 4] = 'e';
            cbuffer_index += 5;
        }
        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_array_union(lua_State* state) {
    NumArray* a = checkarray(state, 1);
    NumArray* b = checkarray(state, 2);

    NumArray* u = array_new(state, fmax(a->size, b->size));

    for (size_t i = 0; i < fmax(N_WORDS(a->size), N_WORDS(b->size)); i++) {
        int value = 0;

        if (i < N_WORDS(a->size)) {
            value |= a->values[i];
        }

        if (i < N_WORDS(b->size)) {
            value |= b->values[i];
        }

        u->values[i] = value;
    }

    return 1;
}

int lua_array_intersection(lua_State* state) {
    NumArray* a = checkarray(state, 1);
    NumArray* b = checkarray(state, 2);

    NumArray* u = array_new(state, fmax(a->size, b->size));

    for (size_t i = 0; i < std::min(N_WORDS(a->size), N_WORDS(b->size)); i++) {
        u->values[i] = a->values[i] & b->values[i];
    }

    return 1;
}

const struct luaL_Reg arrayLib_f[] = {
    {"new", lua_array_new},
    {NULL, NULL}
};

const struct luaL_Reg arrayLib_m[] = {
    {"__tostring", array2string},
    {"set", setarray},
    {"get", getarray},
    {"size", getsize},
    {"union", lua_array_union},
    {"intersection", lua_array_intersection},
    {NULL, NULL}
};

LUALIB_API int luaopen_arrayLib(lua_State* state) {
    luaL_newmetatable(state, "LuaBook.array");

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

    luaL_setfuncs(state, arrayLib_m, 0);

    luaL_newlib(state, arrayLib_f);
    lua_setglobal(state, "array");

    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_arrayLib(L);

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

    return 0;
}

main.lua

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

a:set(1, true)
a:set(5, true)
a:set(7, true)
a:set(8, true)

print("a:", a)

b:set(1, true)
b:set(2, true)
b:set(5, true)

print("b:", b)

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

output.txt

a:	[true, false, false, false, true, false, true, true, false, false]
b:	[true, true, false, false, true]
intersection:	[true, false, false, false, true, false, false, false, false, false]
union:	[true, true, false, false, true, false, true, true, false, false]