1 files changed, 56 insertions, 13 deletions

diff --git a/lbitlib.c b/lbitlib.c
index e3ce2f3c..3e86bab1 100644
--- a/lbitlib.c
+++ b/lbitlib.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lbitlib.c,v 1.12 2010/11/22 16:39:20 roberto Exp roberto $
+** $Id: lbitlib.c,v 1.13 2010/11/22 18:06:33 roberto Exp roberto $
 ** Standard library for bitwise operations
 ** See Copyright Notice in lua.h
 */
@@ -18,21 +18,23 @@
 
 #define ALLONES         (~(((~(lua_Unsigned)0) << (NBITS - 1)) << 1))
 
-/* mask to trim extra bits */
+/* macro to trim extra bits */
 #define trim(x)         ((x) & ALLONES)
 
 
-typedef lua_Unsigned b_uint;
+/* builds a number with 'n' ones (1 <= n <= NBITS) */
+#define mask(n)         (~((ALLONES << 1) << ((n) - 1)))
+
 
+typedef lua_Unsigned b_uint;
 
-#define getuintarg(L,arg)       luaL_checkunsigned(L,arg)
 
 
 static b_uint andaux (lua_State *L) {
   int i, n = lua_gettop(L);
   b_uint r = ~(b_uint)0;
   for (i = 1; i <= n; i++)
-    r &= getuintarg(L, i);
+    r &= luaL_checkunsigned(L, i);
   return trim(r);
 }
 
@@ -55,7 +57,7 @@ static int b_or (lua_State *L) {
   int i, n = lua_gettop(L);
   b_uint r = 0;
   for (i = 1; i <= n; i++)
-    r |= getuintarg(L, i);
+    r |= luaL_checkunsigned(L, i);
   lua_pushunsigned(L, trim(r));
   return 1;
 }
@@ -65,14 +67,14 @@ static int b_xor (lua_State *L) {
   int i, n = lua_gettop(L);
   b_uint r = 0;
   for (i = 1; i <= n; i++)
-    r ^= getuintarg(L, i);
+    r ^= luaL_checkunsigned(L, i);
   lua_pushunsigned(L, trim(r));
   return 1;
 }
 
 
 static int b_not (lua_State *L) {
-  b_uint r = ~getuintarg(L, 1);
+  b_uint r = ~luaL_checkunsigned(L, 1);
   lua_pushunsigned(L, trim(r));
   return 1;
 }
@@ -96,17 +98,17 @@ static int b_shift (lua_State *L, b_uint r, int i) {
 
 
 static int b_lshift (lua_State *L) {
-  return b_shift(L, getuintarg(L, 1), luaL_checkint(L, 2));
+  return b_shift(L, luaL_checkunsigned(L, 1), luaL_checkint(L, 2));
 }
 
 
 static int b_rshift (lua_State *L) {
-  return b_shift(L, getuintarg(L, 1), -luaL_checkint(L, 2));
+  return b_shift(L, luaL_checkunsigned(L, 1), -luaL_checkint(L, 2));
 }
 
 
 static int b_arshift (lua_State *L) {
-  b_uint r = getuintarg(L, 1);
+  b_uint r = luaL_checkunsigned(L, 1);
   int i = luaL_checkint(L, 2);
   if (i < 0 || !(r & ((b_uint)1 << (NBITS - 1))))
     return b_shift(L, r, -i);
@@ -121,7 +123,7 @@ static int b_arshift (lua_State *L) {
 
 
 static int b_rot (lua_State *L, int i) {
-  b_uint r = getuintarg(L, 1);
+  b_uint r = luaL_checkunsigned(L, 1);
   i &= (NBITS - 1);  /* i = i % NBITS */
   r = trim(r);
   r = (r << i) | (r >> (NBITS - i));
@@ -140,17 +142,58 @@ static int b_rrot (lua_State *L) {
 }
 
 
+/*
+** get field and width arguments for field-manipulation functions,
+** checking whether they are valid
+*/
+static int fieldargs (lua_State *L, int farg, int *width) {
+  int f = luaL_checkint(L, farg);
+  int w = luaL_optint(L, farg + 1, 1);
+  luaL_argcheck(L, 0 <= f, farg, "field cannot be netative");
+  luaL_argcheck(L, 0 < w, farg + 1, "width must be positive");
+  if (f + w > NBITS)
+    luaL_error(L, "trying to access non-existent bits");
+  *width = w;
+  return f;
+}
+
+
+static int b_extract (lua_State *L) {
+  int w;
+  b_uint r = luaL_checkunsigned(L, 1);
+  int f = fieldargs(L, 2, &w);
+  r = (r >> f) & mask(w);
+  lua_pushunsigned(L, r);
+  return 1;
+}
+
+
+static int b_replace (lua_State *L) {
+  int w;
+  b_uint r = luaL_checkunsigned(L, 1);
+  b_uint v = luaL_checkunsigned(L, 2);
+  int f = fieldargs(L, 3, &w);
+  int m = mask(w);
+  v &= m;  /* erase bits outside given width */
+  r = (r & ~(m << f)) | (v << f);
+  lua_pushunsigned(L, r);
+  return 1;
+}
+
+
 static const luaL_Reg bitlib[] = {
   {"arshift", b_arshift},
   {"band", b_and},
   {"bnot", b_not},
   {"bor", b_or},
   {"bxor", b_xor},
+  {"btest", b_test},
+  {"extract", b_extract},
   {"lrotate", b_lrot},
   {"lshift", b_lshift},
+  {"replace", b_replace},
   {"rrotate", b_rrot},
   {"rshift", b_rshift},
-  {"btest", b_test},
   {NULL, NULL}
 };