Test/benchmark of Wikipedia Thue-morse implementations

thue-morse.py

# Implementations from Wikipedia
# https://en.wikipedia.org/wiki/Thue%E2%80%93Morse_sequence

from math import ceil, log2
from sys import argv

def thue_morse_fast(n):
    def generate_sequence(seq_length):
        """Thue–Morse sequence."""
        value = 1
        for n in range(seq_length):
            # Note: assumes that (-1).bit_length() gives 1
            x = (n ^ (n - 1)).bit_length() + 1
            if x & 1 == 0:
                # Bit index is even, so toggle value
                value = 1 - value
            yield value

    s = 0
    for b in generate_sequence(n):
        s = (s << 1) | b
    return s

def thue_morse_doubling(n):
    """Return an int containing the first n bits of the Thue-Morse sequence, low-order bit 1st."""
    m = ceil(log2(n))
    bits = 0
    for i in range(m):
        bits |= ((1 << (1 << i)) - 1 - bits) << (1 << i)
    bits >>= 2 ** m - n
    return bits

n = 10 ** int(argv[2])
if argv[1] == "fast":
    thue_morse_fast(n)
elif argv[1] == "doubling":
    thue_morse_doubling(n)
elif argv[1] == "check":
    f = thue_morse_fast(n)
    d = thue_morse_doubling(n)
    assert f == d, f"{f&0xff:0b} ≠ {d&0xff:0b}"
else:
    raise Exception("unknown mode")