Implement a basic bitstream

eos/core/bigint/js.py

@@ -56,8 +56,10 @@ class BigInt(EosObject):
 			('__sub__', 'subtract'),
 			('__mul__', 'multiply'),
 			('__mod__', 'mod'),
+			('__and__', 'and'),
 			('__or__', 'or'),
 			('__lshift__', 'shiftLeft'),
+			('__rshift__', 'shiftRight'),
 			('__xor__', 'xor')
 		]:
 			def make_operator_func(func_):
@@ -98,6 +100,9 @@ class BigInt(EosObject):
 			modulo = BigInt(modulo)
 		return BigInt(self.impl.modPow(other.impl, modulo.impl))
 	
+	def nbits(self):
+		return self.impl.bitLength()
+	
 	def serialise(self):
 		return str(self)
 	

eos/core/bigint/python.py

@@ -16,6 +16,8 @@
 
 from eos.core.objects import EosObject
 
+import math
+
 import random
 system_random = random.SystemRandom()
 
@@ -34,6 +36,9 @@ class BigInt(EosObject):
 			modulo = BigInt(modulo)
 		return BigInt(self.impl.__pow__(other.impl, modulo.impl))
 	
+	def nbits(self):
+		return math.ceil(math.log2(self.impl))
+	
 	def serialise(self):
 		return str(self)
 	
@@ -50,7 +55,7 @@ class BigInt(EosObject):
 	def crypto_random(cls, lower_bound, upper_bound):
 		return cls(system_random.randint(int(lower_bound), int(upper_bound)))
 
-for func in ['__add__', '__sub__', '__mul__', '__mod__', '__or__', '__lshift__', '__xor__']:
+for func in ['__add__', '__sub__', '__mul__', '__mod__', '__and__', '__or__', '__lshift__', '__rshift__', '__xor__']:
 	def make_operator_func(func_):
 		# Create a closure
 		def operator_func(self, other):

eos/core/bitstring.py

@@ -1,17 +0,0 @@
-#   Eos - Verifiable elections
-#   Copyright © 2017  RunasSudo (Yingtong Li)
-#
-#   This program is free software: you can redistribute it and/or modify
-#   it under the terms of the GNU Affero General Public License as published by
-#   the Free Software Foundation, either version 3 of the License, or
-#   (at your option) any later version.
-#
-#   This program is distributed in the hope that it will be useful,
-#   but WITHOUT ANY WARRANTY; without even the implied warranty of
-#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-#   GNU Affero General Public License for more details.
-#
-#   You should have received a copy of the GNU Affero General Public License
-#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
-
-from eos.core.objects import *

eos/core/tests.py

@@ -15,7 +15,6 @@
 #   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from eos.core.bigint import *
-from eos.core.bitstring import *
 from eos.core.objects import *
 
 # Common library things
@@ -103,6 +102,6 @@ class BigIntTestCase(EosTestCase):
 		
 		self.assertEqual(bigint1 + bigint2, 15)
 		self.assertEqual(bigint3 - bigint2, bigint1)
-		self.assertEqual(pow(bigint1, bigint2), pow(5, 10))
-		self.assertEqual(pow(bigint1, bigint2, bigint3), pow(5, 10, 15))
-		self.assertEqual(pow(bigint1, 10, 15), pow(5, 10, 15))
+		self.assertEqual(pow(bigint1, bigint2), 5**10)
+		self.assertEqual(pow(bigint1, bigint2, bigint3), (5**10)%15)
+		self.assertEqual(pow(bigint1, 10, 15), (5**10)%15)

eos/psgjjr/bitstream.py

@@ -0,0 +1,56 @@
+#   Eos - Verifiable elections
+#   Copyright © 2017  RunasSudo (Yingtong Li)
+#
+#   This program is free software: you can redistribute it and/or modify
+#   it under the terms of the GNU Affero General Public License as published by
+#   the Free Software Foundation, either version 3 of the License, or
+#   (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU Affero General Public License for more details.
+#
+#   You should have received a copy of the GNU Affero General Public License
+#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from eos.core.bigint import *
+from eos.core.objects import *
+
+class BitStream(EosObject):
+	def __init__(self, value=None):
+		if value:
+			self.impl = value
+			self.nbits = self.impl.nbits()
+		else:
+			self.impl = ZERO
+			self.nbits = 0
+		self.ptr = 0
+		self.remaining = self.nbits
+	
+	def seek(self, ptr):
+		self.ptr = ptr
+		self.remaining = self.nbits - self.ptr
+	
+	def read(self, nbits):
+		# 11000110110
+		#    ^----
+		val = (self.impl >> (self.remaining - nbits)) & ((ONE << nbits) - ONE)
+		self.ptr += nbits
+		self.remaining -= nbits
+		return val
+	
+	def write(self, bits):
+		# 11     0100110
+		#   10010
+		#   ^----
+		self.impl = ((self.impl >> self.remaining) << (self.remaining + bits.nbits())) | (bits << self.remaining) | (self.impl & ((ONE << self.remaining) - 1))
+		self.ptr += bits.nbits()
+		self.nbits += bits.nbits()
+	
+	def serialise(self):
+		return self.impl
+	
+	@classmethod
+	def deserialise(cls, value):
+		return cls(value)

eos/psgjjr/tests.py

@@ -17,12 +17,33 @@
 from eos.core.tests import *
 
 from eos.core.bigint import *
+from eos.psgjjr.bitstream import *
 from eos.psgjjr.crypto import *
 
 class EGTestCase(EosTestCase):
 	def test_eg(self):
-		pt = BigInt.noncrypto_random(ONE, DEFAULT_GROUP.p - ONE)
+		pt = DEFAULT_GROUP.random_element()
 		sk = EGPrivateKey.generate()
 		ct = sk.public_key.encrypt(pt)
 		m = sk.decrypt(ct)
 		self.assertEqualJSON(pt, m)
+
+class BitStreamTestCase(EosTestCase):
+	def test_bitstream(self):
+		bs = BitStream(BigInt('100101011011', 2))
+		self.assertEqual(bs.read(4), 0b1001)
+		self.assertEqual(bs.read(4), 0b0101)
+		self.assertEqual(bs.read(4), 0b1011)
+		bs = BitStream()
+		bs.write(BigInt('100101011011', 2))
+		bs.seek(0)
+		self.assertEqual(bs.read(4), 0b1001)
+		self.assertEqual(bs.read(4), 0b0101)
+		self.assertEqual(bs.read(4), 0b1011)
+		bs.seek(4)
+		bs.write(BigInt('11', 2))
+		bs.seek(0)
+		self.assertEqual(bs.read(4), 0b1001)
+		self.assertEqual(bs.read(4), 0b1101)
+		self.assertEqual(bs.read(4), 0b0110)
+		self.assertEqual(bs.read(2), 0b11)

eos/tests.py

@@ -17,14 +17,11 @@
 
 from unittest import *
 
-from eos.core.bigint import *
-from eos.core.bitstring import *
-from eos.core.objects import *
-
 import execjs
 
 import importlib
 import os
+import sys
 import types
 
 test_suite = TestSuite()
@@ -70,6 +67,11 @@ for dirpath, dirnames, filenames in os.walk('eos'):
 			obj = getattr(module, name)
 			if isinstance(obj, type):
 				if issubclass(obj, eos.core.tests.EosTestCase):
+					if obj.__module__ != module_name:
+						continue
+					if len(sys.argv) > 1 and not obj.__module__.startswith(sys.argv[1]):
+						continue
+					
 					impl = obj()
 					cls_py = type(name + 'ImplPy', (BasePyTestCase,), {'impl': impl})
 					cls_js = type(name + 'ImplJS', (BaseJSTestCase,), {'module': module_name, 'name': name})

test.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
-coverage run --source=eos --omit='*/js.py','eos/js_tests.py' -m eos.tests || exit 1
+coverage run --source=eos --omit='*/js.py','eos/js_tests.py' -m eos.tests $* || exit 1
 coverage html