bitarray.cpp

/**************************************************************************
 *              Copyright (C), AirM2M Tech. Co., Ltd.
 *
 * Name:    common.cpp
 * Author:  panjun
 * Version: V0.1
 * Date:    2016/10/19
 *
 * Description:
 *     Module that implements a fixed size bit array
 *
 * History:
 *     panjun 16/10/19 Initially create file.
 **************************************************************************/

#include "stdafx.h"
#include "lua.hpp"
#include "platform.h"
#include "platform_conf.h"
#include "auxmods.h"

#define META_NAME                 "eLua.bitarray"
#define bitarray_check( L )      ( bitarray_t* )luaL_checkudata( L, 1, META_NAME )
#define ROUND_SIZE(s)            ( ( ( s ) >> 3 ) + ( ( s ) & 7 ? 1 : 0 ) )

// Unpack modes
enum
{
  BITARRAY_UNPACK_RAW = 0,
  BITARRAY_UNPACK_SEQ
};
 
// Structure that describes our array
typedef struct
{
  UINT32 capacity;
  UINT8 elsize;
  UINT8 values[ 1 ];
} bitarray_t;

// Index shift values/masks
static const UINT8 bitarray_index_shift[] = { 0, 3, 2, 0, 1 };
static const UINT8 bitarray_index_mask[] = { 0, 0x01, 0x03, 0, 0x0F };

// Lua: array = bitarray.new( capacity, [element_size_bits], [fill] ), or
//      array = bitarray.new( "string", [element_size_bits] ), or
//      array = bitarray.new( lua_array, [element_size_bits] )
static int bitarray_new( lua_State *L )
{
  UINT32 total, capacity;
  UINT8 elsize, fill = 0, fromarray = 0;
  const char *buf = NULL;
  bitarray_t *pa;
  size_t temp;
   
  if( lua_isnumber( L, 1 ) )
  {
    // capacity, [element_size_bits], [fill]
    capacity = luaL_checkinteger( L, 1 );
    if( lua_isnumber( L, 2 ) )
      elsize = luaL_checkinteger( L, 2 );
    else
      elsize = 8;
    if( lua_isnumber( L, 3 ) )
      fill = luaL_checkinteger( L, 3 );  
  }
  else if( lua_isstring( L, 1 ) || lua_istable( L, 1 ) /*|| lua_isrotable( L, 1 ) delete by Jack.li ÔÝʱδ¼ÓÈërotable */ )
  {
    // string, [element_size_bits] OR (ro)table, [element_size_bits]
    if( lua_isstring( L, 1 ) )
      buf = lua_tolstring( L, 1, &temp );
    else
    {
      temp = lua_objlen( L, 1 );
      fromarray = 1;
    }
    if( lua_isnumber( L, 2 ) )
      elsize = luaL_checkinteger( L, 2 );
    else
      elsize = 8;
    if( ( temp << 3 ) % elsize )
      return luaL_error( L, "length is not a multiple of element size." );
    capacity = ( temp << 3 ) / elsize;  
  }
  else 
    return luaL_error( L, "invalid arguments." );
    
  if( elsize <= 0 || ( elsize > 32 ) || ( elsize & ( elsize - 1 ) ) )
    return luaL_error( L, "invalid element size." );
  total = ROUND_SIZE( capacity * elsize );
  if( total <= 0 )
    return luaL_error( L, "invalid arguments.");
  pa = ( bitarray_t* )lua_newuserdata( L, sizeof( bitarray_t ) + total - 1 );
  pa->capacity = capacity;
  pa->elsize = elsize;
  
  if( buf )
    memcpy( pa->values, buf, temp );
  else if( fromarray )
  {
    for( total = 1; total <= temp; total ++ )
    {
      lua_rawgeti( L, 1, total );
      pa->values[ total - 1 ] = lua_tointeger( L, -1 );
      lua_pop( L, 1 );
    }
  }
  else
    memset( pa->values, fill, total );
  luaL_getmetatable( L, META_NAME );
  lua_setmetatable( L, -2 );
  return 1;
}

// Helper: get the value at the given index
static UINT32 bitarray_getval( bitarray_t *pa, UINT32 idx )
{
  UINT32 shift, val = 0;
  UINT8 rest, mask;
  
  idx --;
  if( pa->elsize < 8 )        // sub-byte elements
  {
    shift = idx >> bitarray_index_shift[ pa->elsize ];
    mask = 1 << bitarray_index_shift[ pa->elsize ];
    rest = idx & ( mask - 1 );
    val = pa->values[ shift ];
    val = ( val >> ( ( mask - 1 - rest ) * pa->elsize ) ) & bitarray_index_mask[ pa->elsize ]; 
  }
  else      // one byte or more elements
    switch( pa->elsize )
    {
      case 8:
        val = pa->values[ idx ];
        break;
        
      case 16:
        val = *( ( UINT16* )pa->values + idx );
        break;
        
      case 32:
        val = *( ( UINT32* )pa->values + idx );
        break;  
    }
  return val;
}

// Lua: value = array[ idx ]
static int bitarray_get( lua_State *L )
{
  bitarray_t *pa;
  UINT32 idx;
   
  pa = bitarray_check( L );
  idx = luaL_checkinteger( L, 2 );
  if( ( idx <= 0 ) || ( idx > pa->capacity ) )
    return luaL_error( L, "invalid index." );
  lua_pushinteger( L, bitarray_getval( pa, idx ) );    
  return 1;
}

// Lua: array[ key ] = value
static int bitarray_set( lua_State *L )
{
  bitarray_t *pa;
  UINT32 idx, shift, newval, val = 0;
  UINT8 rest, mask;
   
  pa = bitarray_check( L );
  idx = luaL_checkinteger( L, 2 );
  newval = luaL_checkinteger( L, 3 );
  if( ( idx <= 0 ) || ( idx > pa->capacity ) )
    return luaL_error( L, "invalid index." );
  idx --;
  if( pa->elsize < 8 )        // sub-byte elements
  {
    shift = idx >> bitarray_index_shift[ pa->elsize ];
    mask = 1 << bitarray_index_shift[ pa->elsize ];
    rest = idx & ( mask - 1 );
    val = pa->values[ shift ];
    val &= ~( bitarray_index_mask[ pa->elsize ] << ( ( mask - 1 - rest ) * pa->elsize ) );
    val |= newval << ( ( mask - 1 - rest ) * pa->elsize );
    pa->values[ shift ] = val; 
  }
  else      // one byte or more elements
    switch( pa->elsize )
    {
      case 8:
        pa->values[ idx ] = val;;
        break;
        
      case 16:
        *( ( UINT16* )pa->values + idx ) = val;
        break;
        
      case 32:
        *( ( UINT32* )pa->values + idx ) = val;
        break;  
    }    
  return 0;
}

// Lua : size = #array
static int bitarray_len( lua_State *L )
{
  bitarray_t *pa;
  
  pa = bitarray_check( L );
  lua_pushinteger( L, pa->capacity );
  return 1;
}

// Lua iterator
static int bitarray_iter( lua_State *L )
{
  bitarray_t *pa;
  UINT32 idx;
  
  pa = bitarray_check( L );
  idx = luaL_checkinteger( L, 2 ) + 1;
  if( idx <= pa->capacity )
  {
    lua_pushinteger( L, idx );
    lua_pushinteger( L, bitarray_getval( pa, idx ) );
    return 2;
  }
  else
    return 0;
}

// Lua iterator "factory"
static int bitarray_pairs( lua_State *L )
{
#if LUA_OPTIMIZE_MEMORY > 0
  lua_pushlightfunction( L, bitarray_iter );
#else
  lua_pushcclosure( L, bitarray_iter, 0 );
#endif
  lua_pushvalue( L, 1 );
  lua_pushinteger ( L, 0 );
  return 3;
}

// Lua: string = bitarray.tostring( array, ["raw"|"seq"] )
static int bitarray_tostring( lua_State *L )
{
  luaL_Buffer b;
  bitarray_t *pa;
  UINT32 idx;
  UINT8 mode = BITARRAY_UNPACK_SEQ;
  const char *ptextmode;
   
  pa = bitarray_check( L );
  if( lua_isstring( L, 2 ) )
  {
    ptextmode = lua_tostring( L, 2 );
    if( !strcmp( ptextmode, "raw" ) )
      mode = BITARRAY_UNPACK_RAW;
    else if( !strcmp( ptextmode, "seq" ) )
      mode = BITARRAY_UNPACK_SEQ;
    else
      return luaL_error( L, "invalid mode string." );
  }
  if( ( mode == BITARRAY_UNPACK_SEQ ) && ( pa->elsize > 8 ) )
    return luaL_error( L, "element size too large." );
  luaL_buffinit( L, &b );
  if( mode == BITARRAY_UNPACK_SEQ )
    for( idx = 1; idx <= pa->capacity; idx  ++ )
      luaL_addchar( &b, bitarray_getval( pa, idx ) );
  else
    luaL_addlstring( &b, ( char* )pa->values, ROUND_SIZE( pa->capacity * pa->elsize ) );
  luaL_pushresult( &b );
  return 1;  
}

// Lua: table = bitarray.totable( array, ["raw"|"seq"] )
static int bitarray_totable( lua_State *L )
{
  bitarray_t *pa;
  UINT32 idx;
  UINT8 mode = BITARRAY_UNPACK_SEQ;
  const char *ptextmode;
   
  pa = bitarray_check( L );
  if( lua_isstring( L, 2 ) )
  {
    ptextmode = lua_tostring( L, 2 );
    if( !strcmp( ptextmode, "raw" ) )
      mode = BITARRAY_UNPACK_RAW;
    else if( !strcmp( ptextmode, "seq" ) )
      mode = BITARRAY_UNPACK_SEQ;
    else
      return luaL_error( L, "invalid mode string." );
  }
  if( ( mode == BITARRAY_UNPACK_SEQ ) && ( pa->elsize > 8 ) )
    return luaL_error( L, "element size too large." );
  lua_newtable( L );    
  if( mode == BITARRAY_UNPACK_SEQ )
    for( idx = 1; idx <= pa->capacity; idx ++ )
    {
      lua_pushinteger( L, bitarray_getval( pa, idx ) );
      lua_rawseti( L, -2, idx );
    }      
  else
    for( idx = 0; idx < ROUND_SIZE( pa->capacity * pa->elsize ); idx ++ )
    {
      lua_pushinteger( L, pa->values[ idx ] );
      lua_rawseti( L, -2, idx + 1 );
    }
  return 1;  
}

// Module function map
#define MIN_OPT_LEVEL 2
const luaL_Reg bitarray_map[] = 
{
  {"new", bitarray_new},
  {"pairs", bitarray_pairs},
  {"tostring", bitarray_tostring},
  {"totable", bitarray_totable},
  {NULL, NULL} 
};

static const luaL_Reg bitarray_mt_map[] = 
{
  {"__index", bitarray_get},
  {"__newindex", bitarray_set},
  {"__len", bitarray_len},
  {NULL, NULL}
};

LUALIB_API int luaopen_bitarray( lua_State* L )
{
#if LUA_OPTIMIZE_MEMORY > 0
  luaL_rometatable( L, META_NAME, ( void* )bitarray_mt_map );
  return 0;
#else // #if LUA_OPTIMIZE_MEMORY > 0
  luaL_newmetatable( L, META_NAME );
  luaL_register( L, NULL, bitarray_mt_map );
  luaL_register( L, AUXLIB_BITARRAY, bitarray_map );  
  return 1;
#endif // #if LUA_OPTIMIZE_MEMORY > 0  
}