fix mt

Author: maple3142

examples/mt.py

@@ -15,13 +15,14 @@ def timeit(task_name):
         print(f"{task_name} took {end - start:.2f} seconds")
 
 
-def mt19937(bs):
+def mt19937(bs, samples=None):
     print("bs:", bs)
     rand = random.Random(3142)
     st = tuple(rand.getstate()[1][:-1])
 
     effective_bs = ((bs - 1) & bs) or bs
-    out = [rand.getrandbits(bs) for _ in range(624 * 32 // effective_bs)]
+    samples = 624 * 32 // effective_bs if samples is None else samples
+    out = [rand.getrandbits(bs) for _ in range(samples)]
 
     lin = LinearSystem([32] * 624)
     mt = lin.gens()
@@ -39,10 +40,14 @@ def mt19937(bs):
     pyrand = rng.to_python_random()
     assert all(rng.getrandbits(bs) == o for o in out)
     assert all(pyrand.getrandbits(bs) == o for o in out)
+    for _ in range(100):
+        assert rng.getrandbits(bs) == rand.getrandbits(bs)
 
 
 if __name__ == "__main__":
     mt19937(32)
     mt19937(17)
     mt19937(9)
     mt19937(1)
+    mt19937(1337, 19968 // 1337 + 10)
+    mt19937(137, 19968 // 137 + 60)

gf2bv/__init__.py

@@ -52,6 +52,9 @@ def __rshift__(self, n: int):
     def __lshift__(self, n: int):
         return BitVec((0,) * n + self._bits[:-n])
 
+    def lshift_ext(self, n: int):
+        return BitVec((0,) * n + self._bits)
+
     def __and__(self, mask: int):
         bs = to_bits(len(self._bits), mask)
         if all(bs):
@@ -61,7 +64,24 @@ def __and__(self, mask: int):
 
     __rand__ = __and__
 
-    def __or__(self, mask: int):
+    def __or__(self, mask: BitVec | int):
+        if isinstance(mask, BitVec):
+            if len(self._bits) > len(mask._bits):
+                self, mask = mask, self
+            ar = [0] * len(mask._bits)
+            for i in range(len(self._bits)):
+                if self._bits[i] not in (0, 1) and mask._bits[i] not in (0, 1):
+                    raise ValueError(
+                        "Cannot compute logical or using bitvecs with non-zero bits"
+                    )
+                if self._bits[i] == 1 or mask._bits[i] == 1:
+                    ar[i] = 1
+                elif self._bits[i] == 0:
+                    ar[i] = mask._bits[i]
+                elif mask._bits[i] == 0:
+                    ar[i] = self._bits[i]
+            ar[len(self._bits) :] = mask._bits[len(self._bits) :]
+            return BitVec(tuple(ar))
         bs = to_bits(len(self._bits), mask)
         if all(bs):
             # if all bits are set, it becomes all ones

gf2bv/crypto/mt.py

@@ -52,20 +52,29 @@ def __call__(self):
         self.mti += 1
         return self.temper(y)
 
+    def _getrandbits_word(self, k):
+        r = self()
+        if isinstance(r, BitVec):
+            return r[self.w - k :]
+        return r >> (self.w - k)
+
     def getrandbits(self, k=None):
         """Uses the CPython's implementation of random.getrandbits()"""
         if k is None:
             k = self.w
-        if k == 0:
-            return 0
+        if k < 0:
+            raise ValueError("number of bits cannot be negative")
         if k <= self.w:
-            return self.__call__() >> (self.w - k)
+            return self._getrandbits_word(k)
+        words = (k - 1) // self.w + 1
         x = 0
-        for i in range(0, k, self.w):
-            r = self.__call__()
-            if i + self.w > k:
-                r = r >> (self.w - (k - i))
-            x |= r << i
+        for i in range(words):
+            r = self._getrandbits_word(min(k, self.w))
+            if isinstance(r, BitVec):
+                x |= r.lshift_ext(self.w * i)
+            else:
+                x |= r << (self.w * i)
+            k -= self.w
         return x