path: root/bits.py

import re
from operator import __and__
from functools import reduce


class WordBase(object):
    def slice_allowempty(self, range):
        if not isinstance(range, slice):
            range = slice(range, range)
        hi, lo = range.start, range.stop

        if hi < lo:
            return empty

        return self[range]

    def __and__(self, other):
        if isinstance(other, WordBase):
            return BitsUnion(self, other)
        raise NotImplementedError()


class BitsUnion(WordBase):
    def __init__(self, *parts):
        self.parts = parts

    def __repr__(self):
        return "BitsUnion({})".format(", ".join(repr(part) for part in self.parts))

    def as_value(self, **kwargs):
        """reduce down to a single Bitfield value.

        >>> BitsUnion(Bitfield(0x12, 8), Bitfield(0x3, 4), Bitfield(0x4, 4)).as_value()
        Bitfield(0x1234, 16)
        >>> BitsUnion(Bitfield(0x12, 8), LabelRef("test", 0, "imm", slice(11, 4))).as_value(labels={'test': 0xf34f})
        Bitfield(0x1234, 16)
        """
        return reduce(__and__, (part.as_value(**kwargs) for part in self.parts))

    @property
    def size(self):
        """get the total size.

        >>> BitsUnion(Bitfield(0x12, 8), Bitfield(0x3, 4), Bitfield(0x4, 4)).size
        16
        >>> BitsUnion(Bitfield(0x12, 8), LabelRef("test", 0, "imm", slice(11, 4))).size
        16
        """
        return sum(part.size for part in self.parts)

    def __getitem__(self, range):
        """slice a union given hi and lo bit indices.

        >>> BitsUnion(Bitfield(0x12, 8), Bitfield(0x3, 4), Bitfield(0x4, 4))[11:4]
        Bitfield(0x23, 8)
        >>> BitsUnion(Bitfield(0x12, 8), LabelRef("test", 0, "imm", slice(11, 4)))[11:4]
        BitsUnion(Bitfield(0x2, 4), LabelRef('test', 0, ('imm', 32), 11:8))
        """
        if not isinstance(range, slice):
            range = slice(range, range)
        hi, lo = range.start, range.stop

        size = self.size
        if hi < lo:
            raise ValueError("cant slice reverse range")
        elif lo < 0 or hi >= size:
            raise ValueError(
                "slice [{0.start}:{0.stop}] out of range of {1}".format(range, self)
            )

        i = size
        result = Bitfield(0, 0)
        for part in self.parts:
            p_size = part.size
            i -= p_size
            if i > hi:
                # still above hi
                continue

            r_hi = min(p_size - 1, hi - i)
            r_lo = max(0, lo - i)

            result = result & part[r_hi:r_lo]

            if i <= lo:
                # past lo
                break

        return result

    def __and__(self, other):
        if isinstance(other, BitsUnion):
            return BitsUnion(*self.parts, *other.parts)
        elif isinstance(other, WordBase):
            return BitsUnion(*self.parts, other)
        raise NotImplementedError()


class Bitfield(WordBase):
    def __init__(self, val, size, extra=()):
        if val < 0:
            min = -(1 << (size - 1))
            max = -1 - min

            if val > max or val < min:
                raise ValueError(
                    "value {} out of i{} range [{};{}]".format(val, size, min, max)
                )

            if val < 0:
                val = (1 << size) + val

        self.val = val
        self.size = size
        self.extra = extra

    def __repr__(self):
        """format for debugging.
        >>> Bitfield(0, 8)
        Bitfield(0x0, 8)
        >>> Bitfield(16, 8)
        Bitfield(0x10, 8)
        >>> Bitfield(-16, 8)
        Bitfield(0xf0, 8)
        >>> Bitfield(3, 5)
        Bitfield(0x3, 5)
        >>> Bitfield(3, 4, extra=('hello',))
        Bitfield(0x3, 4, extra=('hello',))
        """
        extra = ""
        if self.extra:
            extra = ", extra={!r}".format(self.extra)
        return "Bitfield(0x{:x}, {}{})".format(self.val, self.size, extra)

    def __str__(self):
        """format as a word.
        >>> str(Bitfield(0, 8))
        '00/8'
        >>> str(Bitfield(16, 8))
        '10/8'
        >>> str(Bitfield(3, 8))
        '03/8'
        >>> str(Bitfield(3, 5))
        '03/5'
        >>> str(Bitfield(3, 4))
        '3/4'
        >>> str(Bitfield(3, 4, extra=('hello',)))
        '3/hello/4'
        """
        digits = 1 + ((self.size - 1) // 4)
        base = "{:x}".format(self.val)
        base = (digits - len(base)) * "0" + base
        return "/".join([base, *self.extra, str(self.size)])

    def as_value(self, **kwargs):
        return self

    def __getitem__(self, range):
        """slice a bitfield given hi and lo bit indices.

        >>> Bitfield(0x7f, 8)[7:4]
        Bitfield(0x7, 4)
        >>> Bitfield(0b100, 3)[2:2]
        Bitfield(0x1, 1)
        >>> Bitfield(0x12345678, 32)[7:0]
        Bitfield(0x78, 8)
        >>> Bitfield(0x12345678, 32)[15:8]
        Bitfield(0x56, 8)
        >>> Bitfield(0x12345678, 32)[23:16]
        Bitfield(0x34, 8)

        >>> Bitfield(0xf, 4)[4:0]
        Traceback (most recent call last):
          ...
        ValueError: slice [4:0] out of range of f/4
        >>> Bitfield(0xf, 4)[2:3]
        Traceback (most recent call last):
          ...
        ValueError: cant slice reverse range
        >>> Bitfield(0xf, 4)[3:-1]
        Traceback (most recent call last):
          ...
        ValueError: slice [3:-1] out of range of f/4
        """
        if not isinstance(range, slice):
            range = slice(range, range)
        hi, lo = range.start, range.stop

        if hi < lo:
            raise ValueError("cant slice reverse range")
        elif lo < 0 or hi >= self.size:
            raise ValueError(
                "slice [{0.start}:{0.stop}] out of range of {1}".format(range, self)
            )

        size = hi - lo + 1
        val = (self.val >> lo) & ((1 << size) - 1)
        return Bitfield(val, size)

    def __and__(self, other):
        """concatenate multiple bitfields.

        >>> Bitfield(0b00, 2) & Bitfield(0b10, 2)
        Bitfield(0x2, 4)
        >>> Bitfield(0b110, 3) & Bitfield(0b0110, 4) & Bitfield(0b1, 1)
        Bitfield(0xcd, 8)
        """
        if isinstance(other, Bitfield):
            return Bitfield(self.val << other.size | other.val, self.size + other.size)
        elif isinstance(other, WordBase):
            return BitsUnion(self, other)
        raise NotImplementedError()


global_slice = slice


class LabelRef(WordBase):
    mode_re = re.compile(r"^(imm|off)(\d+)?$")

    def __init__(self, label, offset, mode, slice=None):
        if isinstance(mode, str):
            match = LabelRef.mode_re.match(mode)
            assert match, ValueError("invalid label reference tag")
            mode = (match.group(1), int(match.group(2) or 32))

        self.label = label
        self.mode = mode
        self.slice = slice or global_slice(self.mode[1] - 1, 0)
        self.offset = offset

    def __repr__(self):
        """format for debugging.

        >>> LabelRef("label", 0, ("imm", 32), slice(31,31))
        LabelRef('label', 0, ('imm', 32), 31:31)
        >>> LabelRef("label", 4, ("imm", 32), slice(11,0))
        LabelRef('label', 4, ('imm', 32), 11:0)
        >>> LabelRef("test", -4, "off32")
        LabelRef('test', -4, ('off', 32), 31:0)
        >>> LabelRef("test", 0, "off")
        LabelRef('test', 0, ('off', 32), 31:0)
        """
        return "LabelRef({0!r}, {1}, {2}, {3.start}:{3.stop})".format(
            self.label, self.offset, self.mode, self.slice
        )

    def __str__(self):
        """format as a word.

        >>> str(LabelRef("label", 0, ("imm", 32), slice(31,31)))
        'label/imm32/[31:31]'
        >>> str(LabelRef("label", 0, ("imm", 32), slice(11,0)))
        'label/imm32/[11:0]'
        >>> str(LabelRef("test", 4, "off32"))
        'test+4/off32/[31:0]'
        >>> str(LabelRef("test", -12, "off"))
        'test-12/off32/[31:0]'
        """
        offset = ""
        if self.offset != 0:
            offset = "{:+}".format(self.offset)

        return "{0}{1}/{2[0]}{2[1]}/[{3.start}:{3.stop}]".format(
            self.label, offset, self.mode, self.slice
        )

    @property
    def size(self):
        return self.slice.start - self.slice.stop + 1

    def as_value(self, labels={}, **kwargs):
        """resolve the reference.

        >>> labels = {'test': 0xabcd1234, 'main': 0x8000}
        >>> LabelRef("test", 0, "imm").as_value(labels=labels)
        Bitfield(0xabcd1234, 32)
        >>> LabelRef("test", 0, "imm", slice(15, 0)).as_value(labels=labels)
        Bitfield(0x1234, 16)
        >>> LabelRef("test", 0, "imm", slice(31, 12)).as_value(labels=labels)
        Bitfield(0xabcd1, 20)
        >>> LabelRef("test", 4, "imm", slice(15, 0)).as_value(labels=labels)
        Bitfield(0x1238, 16)
        >>> LabelRef("test", 4, "imm", slice(15, 4)).as_value(labels=labels)
        Bitfield(0x123, 12)
        >>> LabelRef("main", 0, "off").as_value(labels=labels, pc=0x6000)
        Bitfield(0x2000, 32)
        >>> LabelRef("main", 0, "off").as_value(labels=labels, pc=0x800a)
        Bitfield(0xfffffff6, 32)
        >>> LabelRef("main", 0, "off", slice(7, 0)).as_value(labels=labels, pc=0x8012)
        Bitfield(0xee, 8)
        >>> LabelRef("main", 0, "off8").as_value(labels=labels, pc=0x8012)
        Bitfield(0xee, 8)
        >>> LabelRef("main", 0, "off4").as_value(labels=labels, pc=0x8012)
        Traceback (most recent call last):
            ...
        ValueError: value -18 out of i4 range [-8;7]
        >>> LabelRef("main", 4, "off", slice(7, 0)).as_value(labels=labels, pc=0x70fa)
        Bitfield(0xa, 8)
        """
        if self.label not in labels:
            raise ValueError("undefined label '{}'".format(self.label))

        value = labels[self.label] + self.offset

        if self.mode[0] == "off":
            assert "pc" in kwargs, TypeError(
                "'pc' is required to resolve offset-references"
            )
            value = value - kwargs["pc"]

        field = Bitfield(value, self.mode[1])
        return field[self.slice]

    def __getitem__(self, range):
        """slice a label using hi and lo bit indices.

        >>> label = LabelRef("label", 0, "imm32", slice(31,0))
        >>> label.size
        32
        >>> label[31]
        LabelRef('label', 0, ('imm', 32), 31:31)
        >>> label[11:0]
        LabelRef('label', 0, ('imm', 32), 11:0)
        >>> label[11:0].size
        12
        >>> label[12:1]
        LabelRef('label', 0, ('imm', 32), 12:1)
        >>> label = LabelRef("abc", 0, "off32", slice(31,12))
        >>> label.size
        20
        >>> label[9:0]
        LabelRef('abc', 0, ('off', 32), 21:12)
        >>> label[9:2]
        LabelRef('abc', 0, ('off', 32), 21:14)
        >>> label[20:8]
        Traceback (most recent call last):
            ...
        ValueError: slice [20:8] out of range of abc/off32/[31:12]
        """
        if not isinstance(range, slice):
            range = slice(range, range)
        hi, lo = range.start, range.stop

        if hi < lo:
            raise ValueError("cant slice reverse range")
        elif lo < 0 or hi >= self.size:
            raise ValueError(
                "slice [{0.start}:{0.stop}] out of range of {1}".format(range, self)
            )

        new_slice = slice(self.slice.stop + hi, self.slice.stop + lo)
        return LabelRef(self.label, self.offset, self.mode, new_slice)


empty = Bitfield(0, 0)


def u(num, bits):
    """parse an unsigned integer into a bitfield.

    >>> u(0x08, 8)
    Bitfield(0x8, 8)
    >>> u(0xff, 8)
    Bitfield(0xff, 8)

    >>> u(0xf0, 7)
    Traceback (most recent call last):
      ...
    ValueError: value 240 (u8) too large for u7 field
    >>> u(-1, 8)
    Traceback (most recent call last):
      ...
    ValueError: negative value not allowed: -1
    """
    if num < 0:
        raise ValueError("negative value not allowed: {}".format(num))

    if num.bit_length() > bits:
        raise ValueError(
            "value {} (u{}) too large for u{} field".format(num, num.bit_length(), bits)
        )

    return Bitfield(num, bits)


def i(num, bits):
    """parse a signed integer into a bitfield.

    >>> i(8, 8)
    Bitfield(0x8, 8)
    >>> i(-4, 8)
    Bitfield(0xfc, 8)
    >>> i(0x7f, 8)
    Bitfield(0x7f, 8)
    >>> i(-128, 8)
    Bitfield(0x80, 8)

    >>> i(128, 8)
    Traceback (most recent call last):
      ...
    ValueError: value 128 (i9) too large for i8 field [-128;127]
    >>> i(-129, 8)
    Traceback (most recent call last):
      ...
    ValueError: value -129 (i9) too large for i8 field [-128;127]
    """
    min = -(1 << (bits - 1))
    max = -1 - min

    if num > max or num < min:
        raise ValueError(
            "value {} (i{}) too large for i{} field [{};{}]".format(
                num, num.bit_length() + 1, bits, min, max
            )
        )

    if num < 0:
        num = (1 << bits) + num

    return u(num, bits)


def from_part(part):
    """parse a size-tagged subv part into a bit.

    >>> from_part((0x12, 2))
    Bitfield(0x12, 2)
    >>> from_part((0x12, 2, 'extra'))
    Bitfield(0x12, 2)
    """
    val = int(part[0])
    size = part[1]
    return Bitfield(val, int(size))


global_slice = slice
ref_re = re.compile(r"^([^\[+-]+)(?:([+-]\d+))?$")
slice_re = re.compile(r"^\[(\d+):(\d+)]$")


def ref(val, default_slice=None):
    """add a default slice spec to labels if missing.

    >>> ref(('label', 'imm12'))
    LabelRef('label', 0, ('imm', 12), 11:0)
    >>> ref(('label', 'imm12'), slice(12, 1))
    LabelRef('label', 0, ('imm', 12), 12:1)
    >>> ref(('label-4', 'imm12'), slice(12, 1))
    LabelRef('label', -4, ('imm', 12), 12:1)
    >>> ref(('label', 'imm12', '[14:3]'), slice(12, 1))
    LabelRef('label', 0, ('imm', 12), 14:3)
    >>> ref(('label+4', 'imm12', '[14:3]'), slice(12, 1))
    LabelRef('label', 4, ('imm', 12), 14:3)
    >>> ref(('label', 'imm12', '[14:3]'), slice(12, 1))
    LabelRef('label', 0, ('imm', 12), 14:3)
    >>> ref(('label', 'imm12', '[31:0]'), slice(11, 0))
    Traceback (most recent call last):
        ...
    AssertionError: expected 12 bit slice, got label/imm12/[31:0]
    """
    slice = None
    if len(val) == 3:
        label_offset, mode, slice = val
    elif len(val) == 2:
        label_offset, mode = val
    else:
        raise ValueError("expected label reference, got {}".format(val))

    match = ref_re.match(label_offset)
    assert match, ValueError("invalid label reference")
    label, offset = match.group(1), int(match.group(2) or 0)

    if isinstance(slice, str):
        match = slice_re.match(slice)
        assert match, ValueError("invalid slice syntax")
        a, b = match.groups()
        slice = global_slice(int(a), int(b))

    ref = LabelRef(label, offset, mode, slice or default_slice)

    if default_slice:
        expected = default_slice.start - default_slice.stop + 1
        assert expected == ref.size, ValueError(
            "expected {} bit slice, got {}".format(expected, ref)
        )

    return ref