1from __future__ import annotations
2
3import struct
4import math
5import functools
6from typing import Union, Optional, Dict, Callable
7import bitarray
8from bitstring.bitstore import BitStore
9import bitstring
10from bitstring.fp8 import p4binary_fmt, p3binary_fmt
11from bitstring.mxfp import e3m2mxfp_fmt, e2m3mxfp_fmt, e2m1mxfp_fmt, e4m3mxfp_saturate_fmt, e5m2mxfp_saturate_fmt, e4m3mxfp_overflow_fmt, e5m2mxfp_overflow_fmt
12
13# The size of various caches used to improve performance
14CACHE_SIZE = 256
15
16
17def tidy_input_string(s: str) -> str:
18 """Return string made lowercase and with all whitespace and underscores removed."""
19 try:
20 t = s.split()
21 except (AttributeError, TypeError):
22 raise ValueError(f"Expected str object but received a {type(s)} with value {s}.")
23 return ''.join(t).lower().replace('_', '')
24
25
26@functools.lru_cache(CACHE_SIZE)
27def str_to_bitstore(s: str) -> BitStore:
28 _, tokens = bitstring.utils.tokenparser(s)
29 bs = BitStore()
30 for token in tokens:
31 bs += bitstore_from_token(*token)
32 bs.immutable = True
33 return bs
34
35
36def bin2bitstore(binstring: str) -> BitStore:
37 binstring = tidy_input_string(binstring)
38 binstring = binstring.replace('0b', '')
39 try:
40 return BitStore(binstring)
41 except ValueError:
42 raise bitstring.CreationError(f"Invalid character in bin initialiser {binstring}.")
43
44
45def bin2bitstore_unsafe(binstring: str) -> BitStore:
46 return BitStore(binstring)
47
48
49def hex2bitstore(hexstring: str) -> BitStore:
50 hexstring = tidy_input_string(hexstring)
51 hexstring = hexstring.replace('0x', '')
52 try:
53 ba = bitarray.util.hex2ba(hexstring)
54 except ValueError:
55 raise bitstring.CreationError("Invalid symbol in hex initialiser.")
56 return BitStore(ba)
57
58
59def oct2bitstore(octstring: str) -> BitStore:
60 octstring = tidy_input_string(octstring)
61 octstring = octstring.replace('0o', '')
62 try:
63 ba = bitarray.util.base2ba(8, octstring)
64 except ValueError:
65 raise bitstring.CreationError("Invalid symbol in oct initialiser.")
66 return BitStore(ba)
67
68
69def ue2bitstore(i: Union[str, int]) -> BitStore:
70 i = int(i)
71 if i < 0:
72 raise bitstring.CreationError("Cannot use negative initialiser for unsigned exponential-Golomb.")
73 if i == 0:
74 return BitStore('1')
75 tmp = i + 1
76 leadingzeros = -1
77 while tmp > 0:
78 tmp >>= 1
79 leadingzeros += 1
80 remainingpart = i + 1 - (1 << leadingzeros)
81 return BitStore('0' * leadingzeros + '1') + int2bitstore(remainingpart, leadingzeros, False)
82
83
84def se2bitstore(i: Union[str, int]) -> BitStore:
85 i = int(i)
86 if i > 0:
87 u = (i * 2) - 1
88 else:
89 u = -2 * i
90 return ue2bitstore(u)
91
92
93def uie2bitstore(i: Union[str, int]) -> BitStore:
94 i = int(i)
95 if i < 0:
96 raise bitstring.CreationError("Cannot use negative initialiser for unsigned interleaved exponential-Golomb.")
97 return BitStore('1' if i == 0 else '0' + '0'.join(bin(i + 1)[3:]) + '1')
98
99
100def sie2bitstore(i: Union[str, int]) -> BitStore:
101 i = int(i)
102 if i == 0:
103 return BitStore('1')
104 else:
105 return uie2bitstore(abs(i)) + (BitStore('1') if i < 0 else BitStore('0'))
106
107
108def bfloat2bitstore(f: Union[str, float], big_endian: bool) -> BitStore:
109 f = float(f)
110 fmt = '>f' if big_endian else '<f'
111 try:
112 b = struct.pack(fmt, f)
113 except OverflowError:
114 # For consistency we overflow to 'inf'.
115 b = struct.pack(fmt, float('inf') if f > 0 else float('-inf'))
116 return BitStore.frombytes(b[0:2]) if big_endian else BitStore.frombytes(b[2:4])
117
118
119def p4binary2bitstore(f: Union[str, float]) -> BitStore:
120 f = float(f)
121 u = p4binary_fmt.float_to_int8(f)
122 return int2bitstore(u, 8, False)
123
124def p3binary2bitstore(f: Union[str, float]) -> BitStore:
125 f = float(f)
126 u = p3binary_fmt.float_to_int8(f)
127 return int2bitstore(u, 8, False)
128
129def e4m3mxfp2bitstore(f: Union[str, float]) -> BitStore:
130 f = float(f)
131 if bitstring.options.mxfp_overflow == 'saturate':
132 u = e4m3mxfp_saturate_fmt.float_to_int(f)
133 else:
134 u = e4m3mxfp_overflow_fmt.float_to_int(f)
135 return int2bitstore(u, 8, False)
136
137def e5m2mxfp2bitstore(f: Union[str, float]) -> BitStore:
138 f = float(f)
139 if bitstring.options.mxfp_overflow == 'saturate':
140 u = e5m2mxfp_saturate_fmt.float_to_int(f)
141 else:
142 u = e5m2mxfp_overflow_fmt.float_to_int(f)
143 return int2bitstore(u, 8, False)
144
145def e3m2mxfp2bitstore(f: Union[str, float]) -> BitStore:
146 f = float(f)
147 if math.isnan(f):
148 raise ValueError("Cannot convert float('nan') to e3m2mxfp format as it has no representation for it.")
149 u = e3m2mxfp_fmt.float_to_int(f)
150 return int2bitstore(u, 6, False)
151
152def e2m3mxfp2bitstore(f: Union[str, float]) -> BitStore:
153 f = float(f)
154 if math.isnan(f):
155 raise ValueError("Cannot convert float('nan') to e2m3mxfp format as it has no representation for it.")
156 u = e2m3mxfp_fmt.float_to_int(f)
157 return int2bitstore(u, 6, False)
158
159def e2m1mxfp2bitstore(f: Union[str, float]) -> BitStore:
160 f = float(f)
161 if math.isnan(f):
162 raise ValueError("Cannot convert float('nan') to e2m1mxfp format as it has no representation for it.")
163 u = e2m1mxfp_fmt.float_to_int(f)
164 return int2bitstore(u, 4, False)
165
166
167e8m0mxfp_allowed_values = [float(2 ** x) for x in range(-127, 128)]
168def e8m0mxfp2bitstore(f: Union[str, float]) -> BitStore:
169 f = float(f)
170 if math.isnan(f):
171 return BitStore('11111111')
172 try:
173 i = e8m0mxfp_allowed_values.index(f)
174 except ValueError:
175 raise ValueError(f"{f} is not a valid e8m0mxfp value. It must be exactly 2 ** i, for -127 <= i <= 127 or float('nan') as no rounding will be done.")
176 return int2bitstore(i, 8, False)
177
178
179def mxint2bitstore(f: Union[str, float]) -> BitStore:
180 f = float(f)
181 if math.isnan(f):
182 raise ValueError("Cannot convert float('nan') to mxint format as it has no representation for it.")
183 f *= 2 ** 6 # Remove the implicit scaling factor
184 if f > 127: # 1 + 63/64
185 return BitStore('01111111')
186 if f <= -128: # -2
187 return BitStore('10000000')
188 # Want to round to nearest, so move by 0.5 away from zero and round down by converting to int
189 if f >= 0.0:
190 f += 0.5
191 i = int(f)
192 # For ties-round-to-even
193 if f - i == 0.0 and i % 2:
194 i -= 1
195 else:
196 f -= 0.5
197 i = int(f)
198 if f - i == 0.0 and i % 2:
199 i += 1
200 return int2bitstore(i, 8, True)
201
202def int2bitstore(i: int, length: int, signed: bool) -> BitStore:
203 i = int(i)
204 try:
205 x = BitStore(bitarray.util.int2ba(i, length=length, endian='big', signed=signed))
206 except OverflowError as e:
207 if signed:
208 if i >= (1 << (length - 1)) or i < -(1 << (length - 1)):
209 raise bitstring.CreationError(f"{i} is too large a signed integer for a bitstring of length {length}. "
210 f"The allowed range is [{-(1 << (length - 1))}, {(1 << (length - 1)) - 1}].")
211 else:
212 if i >= (1 << length):
213 raise bitstring.CreationError(f"{i} is too large an unsigned integer for a bitstring of length {length}. "
214 f"The allowed range is [0, {(1 << length) - 1}].")
215 if i < 0:
216 raise bitstring.CreationError("uint cannot be initialised with a negative number.")
217 raise e
218 return x
219
220
221def intle2bitstore(i: int, length: int, signed: bool) -> BitStore:
222 x = int2bitstore(i, length, signed).tobytes()
223 return BitStore.frombytes(x[::-1])
224
225
226def float2bitstore(f: Union[str, float], length: int, big_endian: bool) -> BitStore:
227 f = float(f)
228 fmt = {16: '>e', 32: '>f', 64: '>d'}[length] if big_endian else {16: '<e', 32: '<f', 64: '<d'}[length]
229 try:
230 b = struct.pack(fmt, f)
231 except OverflowError:
232 # If float64 doesn't fit it automatically goes to 'inf'. This reproduces that behaviour for other types.
233 b = struct.pack(fmt, float('inf') if f > 0 else float('-inf'))
234 return BitStore.frombytes(b)
235
236
237literal_bit_funcs: Dict[str, Callable[..., BitStore]] = {
238 '0x': hex2bitstore,
239 '0X': hex2bitstore,
240 '0b': bin2bitstore,
241 '0B': bin2bitstore,
242 '0o': oct2bitstore,
243 '0O': oct2bitstore,
244}
245
246
247def bitstore_from_token(name: str, token_length: Optional[int], value: Optional[str]) -> BitStore:
248 if name in literal_bit_funcs:
249 return literal_bit_funcs[name](value)
250 try:
251 d = bitstring.dtypes.Dtype(name, token_length)
252 except ValueError as e:
253 raise bitstring.CreationError(f"Can't parse token: {e}")
254 if value is None and name != 'pad':
255 raise ValueError(f"Token {name} requires a value.")
256 bs = d.build(value)._bitstore
257 if token_length is not None and len(bs) != d.bitlength:
258 raise bitstring.CreationError(f"Token with length {token_length} packed with value of length {len(bs)} "
259 f"({name}:{token_length}={value}).")
260 return bs