enum patch (help request!)

[Asko Kauppi <asko.kauppi@...>]

I got the enum patch working, that's good news.

The bad news is, it's actually slower than my userdata-based approach 
(outside of lua core) used to be..

So.  I'd very much like someone to read the code & comment.  The 
intention is to provide 'class aware' unsigned 32-bit entities that 
greatly help making wrappers for C functions using bitfields, or other 
'magic values'.

By making this a core patch, I was trying to both optimize the 
performance (no mallocs needed) and offer the feature to any Lua users 
(already in LuaX).

Here's the main source file to get you started, for a full Lua 5.1w4 
based tarball, please ask.

-ak

/*
** $Id: lenum.c,v ... $
** Metamethods for enum types
** See Copyright Notice in lua.h
*/

#ifdef ENUM_PATCH

#define lenum_c
#define LUA_CORE

#include "lobject.h"
#include "lstate.h"     // registry()
#include "ltable.h"     // luaH_getnum()
#include "lenum.h"

#include <string.h>
#include <ctype.h>      // tolower()


/*--- Local functions -----------------------------------------*/

// Shift 'val' down to where 'mask' begins
//
static 
lua_Enum Loc_MaskDown( lua_Enum val, lua_Enum mask )
{
    if (mask==0) return 0;  // hmm.. no set bits, weird

    while( (mask&0x01) == 0 )   // shift both mask & value down to 0..N
        { mask >>= 1; val >>= 1; }
    
    return val;
}


/*--- Enumeration methods -----------------------------------------*/

//---
// bool= call( enum a, ["test",] enum b, ... )
// enum= call( enum a, "or"|"and", enum b, ... )
// enum= call( enum a, "xor", [enum b, ...] )
// enum= call( enum a, "<<"|">>" [,int=1] )
// num= call( enum a, "number" )
// str= call( enum a, "string" [,base_uint=10 [,width_uint]] )
//
// Multitude of binary operations can be implemented via the 'call' metamethod.

// All operators have different first characters, this speeds us up.
#define OP_TEST 't'
#define OP_OR   'o'
#define OP_AND  'a'
#define OP_XOR  'x'
#define OP_SHL  '<'
#define OP_SHR  '>'
#define OP_NUM  'n'
#define OP_STR  's'

static
int enum_call( lua_State *L )
{
const char* op= NULL;
char op_c= '\0';
unsigned n=1;
int argn= lua_gettop(L);

    lua_EnumTag tag= lua_enumtag(L,1);
    lua_Enum a= lua_toenum(L,1,tag);

    if (lua_isstring(L,2))
        {
        op= lua_tostring(L,2);
        op_c= tolower(*op);
        n++;
        }

    if ((!op) || (op_c==OP_TEST))    // the only returning 'bool'
        {
        lua_Enum mask=0;
        while( ++n <= argn )
            mask |= lua_toenum(L,n,tag);    // must have same tags

        lua_pushboolean( L, a & mask );
        return 1;   // done! :)
        }

    if (op_c==OP_NUM)
        {
        lua_pushinteger( L, a );
        return 1;
        }

    if (op_c==OP_STR)
        {
        int base= lua_tointeger(L,3);  // 0 == none
        int width= lua_tointeger(L,4);
        char buf[32+1];   // long enough
        
        switch( base )
            {
            case 2: {   // no 'sprintf()' for binary
                char* ptr= buf;
                if (!width) width=32;
                while( width-- )
                    *ptr++= (a & (1<<width)) ? '1':'0';
                *ptr= '\0';
                } break;
            default:    
            case 10:
                sprintf( buf, "%lu", a );    // width ignored
                break;
            case 16:
                sprintf( buf, "%0*lx", width ? width:8, a );
                break;
            }

        lua_pushstring(L,buf);
        return 1;
        }
        
    switch( op_c )  // these all return 'enum'
        {
        case OP_OR:
            while( ++n <= argn )
                a |= lua_toenum(L,n,tag);
            break;

        case OP_AND:
            while( ++n <= argn )
                a &= lua_toenum(L,n,tag);
            break;

        case OP_XOR:
            if (argn==2)
                a ^= 0xffffffffU;
            else
                while( ++n <= argn )
                    a ^= lua_toenum(L,n,tag);
            break;

        case OP_SHL:  // <<
            { int shift= lua_tointeger(L,3);
            a <<= (shift ? shift:1); }
            break;
            
        case OP_SHR:  // >>
            // Note: We're _not_ using arithmetic shift here, although the values
            //       are otherwise regarded as signed. Bit31 always becomes zero.
            //       
            { int shift= lua_tointeger(L,3);
            a= ((unsigned long)a) >> (shift ? shift:1); }
            break;
            
        default:
            lua_pushfstring( L, "Unknown operator: %s", op );
            lua_error( L );
        }

    lua_pushenum( L, a, tag );
    return 1;
}


//---
// int= index( enum a, enum key )
// int= index( enum a, uint bit )
//
// Index notation can be used for extracting a certain flag bit or a masked value
// within the enum. I.e. "var[FLAG]".  Note that unlike the "var(FLAG)" syntax,
// which returns a boolean, we return the actual numeric value, scaled to 0..N.
//
// Use '()' in if/else testing, and [] if actual values are wanted.
//
static
int enum_index( lua_State *L )
{
unsigned long mask;  // must be unsigned, so shift down places 0's at MSB.
int ret=0;

    lua_EnumTag tag= lua_enumtag(L,1);
    lua_Enum a= lua_toenum(L,1,tag);

    if (lua_enumtag(L,2))
        mask= lua_toenum(L,2,tag);     // gives error if wrong family
    else
    if (lua_isinteger(L,2))
        mask= 1L << lua_tointeger(L,2);
    else
        { lua_pushfstring( L, "Bad index: %s", lua_typename(L,2) );
        lua_error( L ); mask=0; }

    ret= Loc_MaskDown( a & mask, mask );
    
    lua_pushinteger( L, ret );
    return 1;
}

//---
// void= newindex( enum a, enum/uint key, enum/uint/bool value )
//
// Setting bits by the table notation:
//
//      myflags[FULLSCREEN]= true
//      myflags[BITFIELD]= 10
//      myflags[BITFIELD]= PRESET_VALUE
//      myflags[20]= 1
//
static
int enum_newindex( lua_State *L )
{
unsigned long mask;
lua_Enum val;
int /*bool*/ shift_up= 1;

    lua_EnumTag tag= lua_enumtag(L,1);
    lua_Enum a= lua_toenum(L,1,tag);

    if (lua_enumtag(L,2))
        mask= lua_toenum(L,2,tag);     // gives error if wrong family
    else
    if (lua_isinteger(L,2))
        mask= 1L << lua_tointeger(L,2);
    else
        { lua_pushfstring( L, "Bad index: %s", lua_typename(L,2) );
        lua_error( L ); mask=0; }

    //---
    if (lua_enumtag(L,3))
        {
        val= lua_toenum(L,3,tag);     // gives error if wrong family
        shift_up= 0;    // already at right positions
        }
    else
    if (lua_isinteger(L,3))
        val= lua_tointeger(L,2);
    else
    if (lua_isnil(L,3) || lua_isboolean(L,3))
        {
        // 'true' sets all the bits of a bitfield (if mask is wider than 1)
        //
        val= lua_toboolean(L,3) ? (lua_Enum)(-1) : 0;
        }
    else
        { lua_pushfstring( L, "Bad value: %s", lua_typename(L,3) );
        lua_error( L ); val=0; }

    //---
    if (shift_up && (val!=0))   // shift according to 'mask'
        {
        unsigned long n= mask;
        while( (n&0x01) == 0 )
            { n >>= 1; val <<= 1; }
        }

    a &= ~mask;         // clear masked bits
    a |= val & mask;    // set new ones

    // In-place modification of the enum value (not public Lua API, but we can.. :)
    //
    // TBD: Or, can we?  Is the value on the stack [1] the real one, or just a copy?
    //
    modevalue( L, 1, a, tag );
    
    return 0;   // void
}


/*--- Enumeration management -----------------------------------------*/

// Some helper macros (from 'gluax.h'):
//
#include <assert.h>
#define /*int abs_index*/ STACK_ABS( L, index ) \
	( (index) >= 0 ? (index) /*absolute*/ : (lua_gettop(L) +(index) +1) )

#define STACK_CHECK(L)     { lua_State* _luastack_= (L); \
                             int _oldtop_= lua_gettop(L);
#define STACK_END(change)  assert( lua_gettop(_luastack_) == _oldtop_+(change) ); }

// Key names used for registry:
//
#define _ENUM       "_ENUM"

//-----
// Initialize the enum storage system
//
void enum_open( lua_State *L )
{
  STACK_CHECK(L) {
    // create class name lookup
    lua_newtable(L);
    lua_setfield(L, LUA_REGISTRYINDEX, _ENUM);

    // Place method closures directly in the registry (as fast access as possible)
    //
    // WARN: This is rather nasty, we use the registry's '__call', '__index' and
    //       '__newindex' names, because that's just easier to do.. (see 'enum_gettm').
    //       Should fix this some day to be more polite. :)
    //
    lua_pushliteral( L, "__call" );   // read by 'G(L)->tmname[TM_CALL]'
    lua_pushcfunction( L, enum_call );
    lua_rawset(L,LUA_REGISTRYINDEX);    // eats both index & function

    lua_pushliteral( L, "__index" );
    lua_pushcfunction( L, enum_index );
    lua_rawset(L,LUA_REGISTRYINDEX);

    lua_pushliteral( L, "__newindex" );
    lua_pushcfunction( L, enum_newindex );
    lua_rawset(L,LUA_REGISTRYINDEX);
    } 
  STACK_END(0);     // should be no change
}

//-----
// Return enum metamethod-like closure
//
const TValue *enum_gettm( lua_State *L, TMS ev )   // TM_CALL, TM_INDEX, TM_NEWINDEX
{
    // WARNING: Not sure if this works.. (we shouldn't touch the stack if possible)
    //
    StkId o= registry(L);    assert(o);
    Table* t= hvalue(o);     assert(t);
    
    const TValue* ret= luaH_getstr( t, G(L)->tmname[ev] );
    
fprintf( stderr, "enum metamethod: %s -> %p\n", 
                    (ev==TM_CALL) ? "TM_CALL" :
                    (ev==TM_INDEX) ? "TM_INDEX" :
                    (ev==TM_NEWINDEX) ? "TM_NEWINDEX" : "", ret );
    assert(ret);
    assert( iscfunction(ret) );
    
    return ret;     // Huh!
}

//-----
// Marry enum 'tt' value and it's id string together. (no divorces allowed!)
// The caller should have checked (should know) that the 'tt' value is free.
//
void enum_marry( lua_State *L, int tt, const char *id )
{
  STACK_CHECK(L) {
    lua_pushliteral( L, _ENUM );
    lua_rawget(L,LUA_REGISTRYINDEX);
    // [-1]= enum lookup table
    
fprintf( stderr, "Marrying %s to %d\n", id, tt );

    lua_pushstring( L, id );    // key
    lua_pushinteger( L, tt );   // value
    lua_rawset(L,-3);           // in the lookup (eats key&val)
    
    lua_pop(L,1);   // remove lookup
    }
  STACK_END(0);   // no change
}

//-----
// Return enum family name
//
const char *enum_name_by_tag( lua_State *L, int tt )
{
const char *ret= NULL;

  STACK_CHECK(L) {
    lua_pushliteral( L, _ENUM );
    lua_rawget(L,LUA_REGISTRYINDEX);
    // [-1]= lookup

    lua_pushnil(L);     // first key (tbl now [-2])
    while (lua_next(L, -2) != 0) {
        // [-3]= lookup
        // [-2]= key (string)
        // [-1]= value (integer)
        //
        if (lua_tointeger(L,-1) == tt) {  // match?
            ret= lua_tostring(L,-2);    // we know it's a string
            lua_pop(L,2);   // removes both value,key (but pointer remains valid until back in Lua)
            break;
            }
        lua_pop(L, 1);  // removes value, keeps key
        }
    // [-1]= lookup
    lua_pop(L,1);
    }
  STACK_END(0);

fprintf( stderr, "enum_name_by_tag: %d -> %s\n", tt, ret?ret:"NULL" );
    
    return ret;
}

//-----
// Return tag if already registered
//
int enum_tag_by_name( lua_State *L, const char *id )
{
int tt= 0;

  STACK_CHECK(L) {
    lua_pushliteral( L, _ENUM );
    lua_rawget(L,LUA_REGISTRYINDEX);
    // [-1]= lookup
    
    lua_pushstring( L, id );    // key
    lua_rawget(L,-2);
    // [-2]= lookup
    // [-1]= value (int) or 'nil'
    
    if (lua_isinteger(L,-1))
        tt= lua_tointeger(L,-1);

    lua_pop(L,2);   // remove pushed value & lookup
    }
  STACK_END(0);

fprintf( stderr, "enum_tag_by_name: %s -> %d\n", id, tt );

    return tt;
}

#endif  // ENUM_PATCH